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

feat(ocr,det,slide): 重构项目结构

Browse Source 
- 优化 规范化模型目录
- 重构 Ocr，Detector，Slide 拆分规范化
This commit is contained in:
CNWei
2026-07-09 19:26:58 +08:00
parent 0cf3d5fefb
commit 2d9cb35590
15 changed files with 431 additions and 608 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
Cargo.toml
View File

@@ -12,6 +12,8 @@ base64 = "0.22.1"
imageproc = { version = "0.26.2", default-features = true }
serde = { version = "1.0.228", features = ["derive"] }
serde_json = "1.0.150"
ndarray="0.16.1"
[features]
default = []

39
src/algo/slide.rs
View File

@@ -1,24 +1,33 @@
use crate::utils::cv_ops;
use crate::utils::cv_ops::{abs_diff, min_max_loc, ndarray_to_luma8, rgb_to_gray};
use crate::utils::image_io::image_to_ndarray;
use anyhow::{Context, Result, anyhow};
use image::{DynamicImage, GenericImageView};
use image::{ImageBuffer, Luma};
use anyhow::{Result, anyhow};
use image::DynamicImage;
use image::Luma;
use imageproc::contrast::{ThresholdType, threshold};
use imageproc::distance_transform::Norm;
use imageproc::edges::canny;
use imageproc::morphology::{close, open};
use imageproc::region_labelling::{Connectivity, connected_components};
use imageproc::template_matching::{MatchTemplateMethod, match_template};
use std::cmp::{max, min};
use tract_onnx::prelude::tract_ndarray::{Array2, Array3, ArrayView2, ArrayView3, Axis, s};
use std::fmt;
use tract_onnx::prelude::tract_ndarray::{ArrayView2, ArrayView3};
#[derive(Debug)]
pub struct SlideResult {
pub target: [i32; 2],
pub target_x: i32,
pub target_y: i32,
pub confidence: f64,
}
impl fmt::Display for SlideResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "滑块匹配测试结果:")?;
writeln!(f, "检测坐标: [x: {}, y: {}]", self.target_x, self.target_y)?;
// 注意：这里保留 4 位小数，如果想让外部控制，也可以直接写 {:.4}
write!(f, "置信度: {:.4}", self.confidence)?;
Ok(())
}
}
pub struct Slider;
@@ -58,23 +67,8 @@ impl Slider {
target: ArrayView3<u8>,
background: ArrayView3<u8>,
) -> Result<SlideResult> {
// let (h, w, _) = target.dim();
// 1. 计算图像差异并灰度化 (对应 cv2.absdiff + cv2.cvtColor)
// 使用 OpenCV 标准权重公式：0.299R + 0.587G + 0.114B
// let mut diff_buffer = ImageBuffer::new(w as u32, h as u32);
// for y in 0..h {
// for x in 0..w {
// let r_diff = (target[[y, x, 0]] as i16 - background[[y, x, 0]] as i16).abs() as f32;
// let g_diff = (target[[y, x, 1]] as i16 - background[[y, x, 1]] as i16).abs() as f32;
// let b_diff = (target[[y, x, 2]] as i16 - background[[y, x, 2]] as i16).abs() as f32;
//
// let gray_diff = (0.299 * r_diff + 0.587 * g_diff + 0.114 * b_diff) as u8;
// diff_buffer.put_pixel(x as u32, y as u32, Luma([gray_diff]));
// }
// }
// 1. 计算差异数组 (复用 cv2::absdiff)
let (th, tw, tc) = target.dim();
let (bh, bw, bc) = background.dim();
@@ -193,9 +187,6 @@ impl Slider {
background: ArrayView2<u8>,
) -> Result<SlideResult> {
// 1. 将 ndarray 转换为 imageproc 需要的 ImageBuffer (无拷贝或轻量转换)
// let (bh, bw) = background.dim();
// 转换逻辑 (假设你已经有方法转回 ImageBuffer)
let t_buf = ndarray_to_luma8(target);
let b_buf = ndarray_to_luma8(background);

170
src/lib.rs
View File

@@ -3,175 +3,7 @@ mod error;
pub mod models;
pub mod utils;
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 (使用内置路径)
// 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(ocr: {})", 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)]
// struct ClassificationBuilder {
// img: DynamicImage,
// png_fix: bool,
// color_filter_colors: Option<Vec<ColorRange>>,
// color_filter_custom_ranges: Option<Vec<ColorRange>>,
// }
// impl ClassificationBuilder {
// pub fn new(img: DynamicImage) -> Self {
// ClassificationBuilder {
// img,
// png_fix: false,
// color_filter_colors: None,
// color_filter_custom_ranges: None,
// }
// }
// pub fn png_fix(mut self, value: bool) -> Self {
// self.png_fix = value;
// self
// }
// pub fn color_filter_colors(mut self, value: Vec<ColorRange>) -> Self {
// self.color_filter_colors = Some(value);
// self
// }
// pub fn color_filter_custom_ranges(mut self, value: Vec<ColorRange>) -> Self {
// self.color_filter_custom_ranges = Some(value);
// self
// }
// pub fn build(self) -> Classification {
// Classification {}
// }
// }
#[cfg(test)]
mod tests {
#[test]
fn test_ctc_decode_indices() {
// 模拟一个 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
// 预期结果索引应该是 [1, 1, 2, 2] 对应的字符
// 具体的断言取决于你的 CHARSET_BETA
// let result = dddd.ctc_decode_indices(&input);
// assert_eq!(result, "AABB");
}
}
pub use models::ocr::metadata::ModelMetadata;

31
src/models/det/executor.rs
View File

@@ -1,15 +1,10 @@
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};
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};
use anyhow::{Context, Result};
use image::{imageops::FilterType, DynamicImage, GenericImageView};
use std::fmt;
use tract_onnx::prelude::tract_ndarray::{prelude::*, s, Array2, Array3, Array4, Axis};
use tract_onnx::prelude::{Tensor};
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)]
@@ -22,15 +17,23 @@ pub struct DetectionResult {
pub class_id: u32,
}
impl fmt::Display for DetectionResult {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
// 结构体只管自己这一行怎么显示，不用管外部的索引 [i]
write!(
f,
"x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
self.x1, self.y1, self.x2, self.y2, self.score, self.class_id
)
}
}
#[derive(Debug)]
pub struct Detector<'a> {
pub(crate) session: &'a DetSession,
#[allow(dead_code)]
pub(crate) use_gpu: bool,
#[allow(dead_code)]
pub(crate) device_id: u8,
}

70
src/models/loader.rs
View File

@@ -1,17 +1,10 @@
use anyhow::{anyhow, Context};
use image::DynamicImage;
use anyhow::Context;
use std::io::Cursor;
use tract_onnx::onnx;
use tract_onnx::prelude::*;
// 关键点：直接使用 tract 重导出的 ndarray
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,
@@ -53,60 +46,3 @@ impl ModelLoader {
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))
// }
// }

3
src/models/ocr/builder.rs
View File

@@ -1,7 +1,8 @@
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;
use crate::models::ocr::token_filter::TokenFilter;
pub struct OcrBuilder {
/// 是否修复PNG格式问题
png_fix: bool,

94
src/models/ocr/color_filter.rs
View File

@@ -1,25 +1,24 @@
use std::str::FromStr;
use crate::utils::cv_ops::rgb_to_opencv_hsv;
use anyhow::anyhow;
use image::{DynamicImage, ImageBuffer, Rgb};
use crate::utils::cv_ops::rgb_to_opencv_hsv;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct HsvRange {
pub lower: (u8, u8, u8), // (H, S, V)
pub upper: (u8, u8, u8), // (H, S, V)
}
use std::str::FromStr;
/// 核心区间判定辅助函数
#[inline(always)]
fn is_pixel_matched(ranges: &[HsvRange], h: u8, s: u8, v: u8) -> bool {
ranges.iter().any(|range| {
h >= range.lower.0 && h <= range.upper.0 &&
s >= range.lower.1 && s <= range.upper.1 &&
v >= range.lower.2 && v <= range.upper.2
h >= range.lower.0
&& h <= range.upper.0
&& s >= range.lower.1
&& s <= range.upper.1
&& v >= range.lower.2
&& v <= range.upper.2
})
}
pub fn filter_image(image: &DynamicImage, hsv_ranges: &[HsvRange]) -> anyhow::Result<DynamicImage> {
pub fn apply_to_image(
image: &DynamicImage,
hsv_ranges: &[HsvRange],
) -> anyhow::Result<DynamicImage> {
// 1. 统一转换为连续内存的 RGB8 缓冲区 (对应 Python 的 Image 到 RGB/BGR 数组转换)
let rgb_img = image.to_rgb8();
let (width, height) = rgb_img.dimensions();
@@ -50,6 +49,13 @@ pub fn filter_image(image: &DynamicImage, hsv_ranges: &[HsvRange]) -> anyhow::Re
Ok(DynamicImage::ImageRgb8(filtered_buffer))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct HsvRange {
pub lower: (u8, u8, u8), // (H, S, V)
pub upper: (u8, u8, u8), // (H, S, V)
}
impl HsvRange {
pub const fn new(lower: (u8, u8, u8), upper: (u8, u8, u8)) -> Self {
Self { lower, upper }
@@ -65,7 +71,8 @@ impl HsvRange {
}
// 2. 校验下界不能大于上界
if self.lower.0 > self.upper.0 || self.lower.1 > self.upper.1 || self.lower.2 > self.upper.2 {
if self.lower.0 > self.upper.0 || self.lower.1 > self.upper.1 || self.lower.2 > self.upper.2
{
return Err("HSV范围下界不能大于上界".to_string());
}
@@ -87,25 +94,58 @@ pub enum ColorPreset {
Custom(Vec<HsvRange>),
}
impl ColorPreset {
impl ColorPreset {
/// 纯裸数据定义，没有任何结构体包装，干净利落
/// 返回值：(范围数量, 范围数组)
/// 完美的零成本抽象：利用常量提升将数据直接打入只读数据段 (.rodata)
pub fn matches(&self) -> &[HsvRange] {
match self {
ColorPreset::Red => &[
HsvRange { lower: (0, 50, 50), upper: (10, 255, 255) },
HsvRange { lower: (170, 50, 50), upper: (180, 255, 255) },
HsvRange {
lower: (0, 50, 50),
upper: (10, 255, 255),
},
HsvRange {
lower: (170, 50, 50),
upper: (180, 255, 255),
},
],
ColorPreset::Blue => &[HsvRange { lower: (100, 50, 50), upper: (130, 255, 255) }],
ColorPreset::Green => &[HsvRange { lower: (40, 50, 50), upper: (80, 255, 255) }],
ColorPreset::Yellow => &[HsvRange { lower: (20, 50, 50), upper: (40, 255, 255) }],
ColorPreset::Orange => &[HsvRange { lower: (10, 50, 50), upper: (20, 255, 255) }],
ColorPreset::Purple => &[HsvRange { lower: (130, 50, 50), upper: (170, 255, 255) }],
ColorPreset::Cyan => &[HsvRange { lower: (80, 50, 50), upper: (100, 255, 255) }],
ColorPreset::Black => &[HsvRange { lower: (0, 0, 0), upper: (180, 255, 50) }],
ColorPreset::White => &[HsvRange { lower: (0, 0, 200), upper: (180, 30, 255) }],
ColorPreset::Gray => &[HsvRange { lower: (0, 0, 50), upper: (180, 30, 200) }],
ColorPreset::Blue => &[HsvRange {
lower: (100, 50, 50),
upper: (130, 255, 255),
}],
ColorPreset::Green => &[HsvRange {
lower: (40, 50, 50),
upper: (80, 255, 255),
}],
ColorPreset::Yellow => &[HsvRange {
lower: (20, 50, 50),
upper: (40, 255, 255),
}],
ColorPreset::Orange => &[HsvRange {
lower: (10, 50, 50),
upper: (20, 255, 255),
}],
ColorPreset::Purple => &[HsvRange {
lower: (130, 50, 50),
upper: (170, 255, 255),
}],
ColorPreset::Cyan => &[HsvRange {
lower: (80, 50, 50),
upper: (100, 255, 255),
}],
ColorPreset::Black => &[HsvRange {
lower: (0, 0, 0),
upper: (180, 255, 50),
}],
ColorPreset::White => &[HsvRange {
lower: (0, 0, 200),
upper: (180, 30, 255),
}],
ColorPreset::Gray => &[HsvRange {
lower: (0, 0, 50),
upper: (180, 30, 200),
}],
ColorPreset::Custom(ranges) => ranges,
}
}
@@ -204,7 +244,6 @@ impl ColorFilter for ColorPreset {
}
}
/// 多路颜色“或”逻辑组合子（并集网络）
pub struct MultiOrColorRestrict<'a> {
pub filters: Vec<&'a dyn ColorFilter>,
@@ -246,4 +285,3 @@ macro_rules! color_any_of {
}
};
}

6
src/models/ocr/executor.rs
View File

@@ -1,7 +1,7 @@
use crate::models::ocr::model_metadata::Resize;
use crate::models::ocr::metadata::Resize;
use crate::models::ocr::session::OcrSession;
use crate::models::ocr::color_filter::{HsvRange, filter_image};
use crate::models::ocr::color_filter::{HsvRange, apply_to_image};
use crate::utils::image_io::png_rgba_white_preprocess;
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
use anyhow::Result;
@@ -141,7 +141,7 @@ impl<'a> Ocr<'a> {
}
Ok(Some(ranges)) => {
// 只有真正需要过滤时，才在内部提取像素并生成清洗后的 Owned 新图
let filtered_img = filter_image(image, ranges)?;
let filtered_img = apply_to_image(image, ranges)?;
Cow::Owned(filtered_img)
}
};

145
src/models/ocr/model_metadata.rs → src/models/ocr/metadata.rs
View File

@@ -1,11 +1,72 @@
use crate::models::ocr::charset::{CHARSET_BETA, CHARSET_OLD, Charset};
use anyhow::{Result, anyhow};
use anyhow::{anyhow, Result};
use serde::Deserialize;
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::Read;
use std::path::Path;
use std::collections::HashMap;
// ==========================================
// 3. 字符集核心结构体 (重命名为 Charset)
// ==========================================
#[derive(Debug, Clone)]
pub struct Charset {
// 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
pub tokens: Vec<Cow<'static, str>>,
// 反向查找表，保证字符转索引为 O(1)
pub char_to_idx: HashMap<Cow<'static, str>, usize>,
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
// pub valid_indices: HashSet<usize>,
}
impl Charset {
// 内部底层统一收拢构造
pub fn new(tokens: Vec<Cow<'static, str>>) -> Self {
let mut char_to_idx = HashMap::with_capacity(tokens.len());
for (idx, token) in tokens.iter().enumerate() {
char_to_idx.entry(token.clone()).or_insert(idx);
// 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑)
// char_to_idx.entry(token.to_string()).or_insert(idx);
}
Self {
tokens,
char_to_idx,
}
}
// --- 业务策略方法 ---
/// 将字符转为索引，不存在返回 -1 (保持与原 Python 库行为一致)
pub fn char_to_index(&self, char_str: &str) -> i32 {
if let Some(&idx) = self.char_to_idx.get(char_str) {
idx as i32
} else {
-1
}
}
/// 将索引转为字符引用，零拷贝。若越界返回 None
pub fn index_to_char_ref(&self, index: usize) -> Option<&str> {
self.tokens.get(index).map(|cow| cow.as_ref())
}
pub fn is_valid_char(&self, char_str: &str) -> bool {
self.char_to_idx.get(char_str).is_some()
}
pub fn size(&self) -> usize {
self.tokens.len()
}
}
// ==========================================
// 4. 标准 Display 接口实现 (对应 __str__)
// ==========================================
impl std::fmt::Display for Charset {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Charset [Total Size: {}", self.size(),)
}
}
// =====================================================================
// 1. 辅助定义的枚举与结构体
// =====================================================================
@@ -74,28 +135,6 @@ pub struct ModelMetadata {
impl ModelMetadata {
// --- 优雅的工厂模式构造器 ---
/// 从预设的旧版字符集创建
pub fn from_builtin_old() -> Self {
Self::from_static_slice(
CHARSET_OLD,
false,
Resize::DynamicWidth(64),
1,
Normalization::ZeroToOne,
)
}
/// 从预设的 Beta 版字符集创建
pub fn from_builtin_beta() -> Self {
Self::from_static_slice(
CHARSET_BETA,
false,
Resize::DynamicWidth(64),
1,
Normalization::MinusOneToOne,
)
}
/// 通用的静态切片转换构造器
pub fn from_static_slice(
slice: &[&'static str],
@@ -158,54 +197,4 @@ impl ModelMetadata {
.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() {
return Err(anyhow!("模型元数据配置文件不存在: {:?}", path));
}
let mut file = File::open(path)?;
let mut content = String::new();
file.read_to_string(&mut content)?;
let dto: ModelMetadataDto = serde_json::from_str(&content)
.map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?;
// 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,
})
}
}

4
src/models/ocr/mod.rs
View File

@@ -1,9 +1,9 @@
mod builder;
mod executor;
mod session;
pub mod charset;
pub mod model_metadata;
pub mod metadata;
pub mod color_filter;
mod token_filter;
pub use builder::OcrBuilder;
pub use executor::{Ocr, OcrResult};

2
src/models/ocr/session.rs
View File

@@ -1,4 +1,4 @@
use crate::models::ocr::model_metadata::ModelMetadata;
use crate::models::ocr::metadata::ModelMetadata;
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
use anyhow::Context;
use anyhow::Result;

146
src/models/ocr/token_filter.rs Normal file
View File

@@ -0,0 +1,146 @@
use std::borrow::Cow;
/// 字符集范围限制枚举
pub struct ValidationCtx<'a> {
pub text: &'a str, // 当前 Token 的文本内容
pub token_id: usize, // 当前 Token 的 ID 索引
}
/// 统一的约束接口
pub trait TokenFilter {
fn matches(&self, ctx: &ValidationCtx) -> bool;
/// 预估容量提示，帮助精准开辟 Vec 内存
fn estimated_capacity(&self) -> usize {
128
}
/// 【新引入的架构级核心方法】
/// 统一接管全量字符集的密集遍历、CTC Blank放行、去重、排序及空交集退化兜底
fn apply_to_charset(&self, tokens: &[Cow<str>]) -> Option<Vec<usize>> {
let mut has_any_match = false;
let estimated_capacity = self.estimated_capacity();
// 1. 精准开辟内存，完美利用容量提示，避免动态乱涨
let mut temp_indices = Vec::with_capacity(estimated_capacity.max(16));
// 2. 高性能原地单次流式迭代
for (idx, token) in tokens.iter().enumerate() {
let token_str = token.as_ref();
// 规则 A: CTC Blank 空字符串或 0 号索引无条件放行
if token_str.is_empty() || idx == 0 {
temp_indices.push(idx);
continue; // 关键：直接跳过，防止后续 matches 匹配成功导致重复 push 产生 Bug
}
// 规则 B: 组装无拷贝上下文
let ctx = ValidationCtx {
text: token_str,
token_id: idx,
};
// 规则 C: 路由到各自具体实现的特异性匹配中（如 Digit 判定、TopN 判定、组合子判定等）
if self.matches(&ctx) {
temp_indices.push(idx);
has_any_match = true;
}
}
// 3. 终极防御：如果整个模型字符集除了 Blank，一个都没对上，直接退化为 None（全量识别）
if !has_any_match {
println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。");
None
} else {
// 4. 排序并去重，为 Ocr 引擎后续进行极其高频的『二分查找』筑起绝对安全的底层保障
temp_indices.sort_unstable();
temp_indices.dedup();
Some(temp_indices)
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CharRestrict {
Digit,
Lowercase,
Uppercase,
CustomList(Vec<String>),
}
impl TokenFilter for CharRestrict {
fn matches(&self, ctx: &ValidationCtx) -> bool {
match self {
Self::Digit => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_digit(),
Self::Lowercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_lowercase(),
Self::Uppercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_uppercase(),
Self::CustomList(vec) => vec.iter().any(|t| t == ctx.text),
}
}
fn estimated_capacity(&self) -> usize {
match self {
Self::Digit => 16,
Self::Lowercase | Self::Uppercase => 32,
Self::CustomList(vec) => vec.len() + 1,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IdRestrict {
TopN(usize),
IdRange(std::ops::Range<usize>),
IdList(Vec<usize>),
}
impl TokenFilter for IdRestrict {
fn matches(&self, ctx: &ValidationCtx) -> bool {
match self {
Self::TopN(n) => ctx.token_id < *n,
Self::IdRange(range) => range.contains(&ctx.token_id),
Self::IdList(vec) => vec.contains(&ctx.token_id),
}
}
fn estimated_capacity(&self) -> usize {
match self {
Self::TopN(n) => *n + 1,
// 2. IdRange：标准标准库 Range 的长度
// 注意：因为范围可能是 1000..2000，它的 len() 返回的是 usize
Self::IdRange(range) => range.len() + 1,
// 3. IdList：Vec 里的元素个数
Self::IdList(vec) => vec.len() + 1,
}
}
}
/// 多路“或”逻辑组合子（支持 N 个规则无缝并集）
pub struct MultiOrRestrict<'a> {
pub filters: Vec<&'a dyn TokenFilter>,
}
impl<'a> TokenFilter for MultiOrRestrict<'a> {
fn matches(&self, ctx: &ValidationCtx) -> bool {
// 核心高阶函数：只要有一个过滤器命中，该 Token 即可放行
self.filters.iter().any(|f| f.matches(ctx))
}
fn estimated_capacity(&self) -> usize {
// 将所有过滤器的预估容量累加，作为最终容量参考
self.filters.iter().map(|f| f.estimated_capacity()).sum()
}
}
// =====================================================================
// 声明式宏：替代 `+` 运算符，解决组合扩展痛苦
// =====================================================================
#[macro_export]
macro_rules! any_of {
// 场景 A：如果用户只传了一个规则，免去构建 Vec 的开销，直接返回其引用
($only:expr) => {
&$only as &dyn $crate::TokenFilter
};
// 场景 B：如果用户传入了多个规则，自动织成一张静态组合网
($($filter:expr),+ $(,)?) => {
&$crate::MultiOrRestrict {
filters: vec![ $( &$filter as &dyn $crate::TokenFilter ),+ ]
}
};
}

9
src/utils/image_processor.rs
View File

@@ -1,5 +1,7 @@
use image::{DynamicImage, GrayImage, imageops::FilterType};
use anyhow::Result;
use image::{DynamicImage, GrayImage, imageops::FilterType, Rgb, ImageBuffer};
use anyhow::{anyhow, Result};
use crate::models::ocr::color_filter::HsvRange;
use crate::utils::cv_ops::rgb_to_opencv_hsv;
/// 对应 Python 的 convert_to_grayscale
/// 将图像转换为灰度图 (L模式)
@@ -34,4 +36,5 @@ pub fn resize_image(
// target_height,
// FilterType::Lanczos3
// )
// }
// }

280
src/models/ocr/charset.rs → tests/char_slice.rs
View File

@@ -1,3 +1,10 @@
use std::borrow::Cow;
use std::fs::File;
use std::path::Path;
use anyhow::anyhow;
use ddddocr_rs::models::ocr::metadata::Charset;
use ddddocr_rs::models::ocr::metadata::{Normalization, Resize};
pub const CHARSET_BETA: &[&str] = &[
"", "笤", "谴", "膀", "荔", "佰", "电", "臁", "矍", "同", "奇", "芄", "吠", "6", "曛", "荇",
"砥", "蹅", "晃", "厄", "殣", "ｃ", "辱", "钋", "杻", "價", "眙", "鴿", "⒄", "裙", "训", "涛",
@@ -517,212 +524,77 @@ pub const CHARSET_BETA: &[&str] = &[
pub const CHARSET_OLD: &[&str] = &["", "笤", "谴", "膀", "荔"];
use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
/// 字符集范围限制枚举
pub struct ValidationCtx<'a> {
pub text: &'a str, // 当前 Token 的文本内容
pub token_id: usize, // 当前 Token 的 ID 索引
}
/// 统一的约束接口
pub trait TokenFilter {
fn matches(&self, ctx: &ValidationCtx) -> bool;
/// 预估容量提示，帮助精准开辟 Vec 内存
fn estimated_capacity(&self) -> usize {
128
}
/// 【新引入的架构级核心方法】
/// 统一接管全量字符集的密集遍历、CTC Blank放行、去重、排序及空交集退化兜底
fn apply_to_charset(&self, tokens: &[Cow<str>]) -> Option<Vec<usize>> {
let mut has_any_match = false;
let estimated_capacity = self.estimated_capacity();
// 1. 精准开辟内存，完美利用容量提示，避免动态乱涨
let mut temp_indices = Vec::with_capacity(estimated_capacity.max(16));
// pub fn from_builtin_old() -> Self {
// Self::from_static_slice(
// CHARSET_OLD,
// false,
// Resize::DynamicWidth(64),
// 1,
// Normalization::ZeroToOne,
// )
// }
//
// /// 从预设的 Beta 版字符集创建
// pub fn from_builtin_beta() -> Self {
// Self::from_static_slice(
// CHARSET_BETA,
// false,
// Resize::DynamicWidth(64),
// 1,
// Normalization::MinusOneToOne,
// )
// }
// 2. 高性能原地单次流式迭代
for (idx, token) in tokens.iter().enumerate() {
let token_str = token.as_ref();
// 规则 A: CTC Blank 空字符串或 0 号索引无条件放行
if token_str.is_empty() || idx == 0 {
temp_indices.push(idx);
continue; // 关键：直接跳过，防止后续 matches 匹配成功导致重复 push 产生 Bug
}
// 规则 B: 组装无拷贝上下文
let ctx = ValidationCtx {
text: token_str,
token_id: idx,
};
// 规则 C: 路由到各自具体实现的特异性匹配中（如 Digit 判定、TopN 判定、组合子判定等）
if self.matches(&ctx) {
temp_indices.push(idx);
has_any_match = true;
}
}
// 3. 终极防御：如果整个模型字符集除了 Blank，一个都没对上，直接退化为 None（全量识别）
if !has_any_match {
println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。");
None
} else {
// 4. 排序并去重，为 Ocr 引擎后续进行极其高频的『二分查找』筑起绝对安全的底层保障
temp_indices.sort_unstable();
temp_indices.dedup();
Some(temp_indices)
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum CharRestrict {
Digit,
Lowercase,
Uppercase,
CustomList(Vec<String>),
}
impl TokenFilter for CharRestrict {
fn matches(&self, ctx: &ValidationCtx) -> bool {
match self {
Self::Digit => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_digit(),
Self::Lowercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_lowercase(),
Self::Uppercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_uppercase(),
Self::CustomList(vec) => vec.iter().any(|t| t == ctx.text),
}
}
fn estimated_capacity(&self) -> usize {
match self {
Self::Digit => 16,
Self::Lowercase | Self::Uppercase => 32,
Self::CustomList(vec) => vec.len() + 1,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum IdRestrict {
TopN(usize),
IdRange(std::ops::Range<usize>),
IdList(Vec<usize>),
}
impl TokenFilter for IdRestrict {
fn matches(&self, ctx: &ValidationCtx) -> bool {
match self {
Self::TopN(n) => ctx.token_id < *n,
Self::IdRange(range) => range.contains(&ctx.token_id),
Self::IdList(vec) => vec.contains(&ctx.token_id),
}
}
fn estimated_capacity(&self) -> usize {
match self {
Self::TopN(n) => *n + 1,
// 2. IdRange：标准标准库 Range 的长度
// 注意：因为范围可能是 1000..2000，它的 len() 返回的是 usize
Self::IdRange(range) => range.len() + 1,
// 3. IdList：Vec 里的元素个数
Self::IdList(vec) => vec.len() + 1,
}
}
}
/// 多路“或”逻辑组合子（支持 N 个规则无缝并集）
pub struct MultiOrRestrict<'a> {
pub filters: Vec<&'a dyn TokenFilter>,
}
impl<'a> TokenFilter for MultiOrRestrict<'a> {
fn matches(&self, ctx: &ValidationCtx) -> bool {
// 核心高阶函数：只要有一个过滤器命中，该 Token 即可放行
self.filters.iter().any(|f| f.matches(ctx))
}
fn estimated_capacity(&self) -> usize {
// 将所有过滤器的预估容量累加，作为最终容量参考
self.filters.iter().map(|f| f.estimated_capacity()).sum()
}
}
// =====================================================================
// 声明式宏：替代 `+` 运算符，解决组合扩展痛苦
// =====================================================================
#[macro_export]
macro_rules! any_of {
// 场景 A：如果用户只传了一个规则，免去构建 Vec 的开销，直接返回其引用
($only:expr) => {
&$only as &dyn $crate::TokenFilter
};
// 场景 B：如果用户传入了多个规则，自动织成一张静态组合网
($($filter:expr),+ $(,)?) => {
&$crate::MultiOrRestrict {
filters: vec![ $( &$filter as &dyn $crate::TokenFilter ),+ ]
}
};
}
// ==========================================
// 3. 字符集核心结构体 (重命名为 Charset)
// ==========================================
#[derive(Debug, Clone)]
pub struct Charset {
// 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
pub tokens: Vec<Cow<'static, str>>,
// 反向查找表，保证字符转索引为 O(1)
pub char_to_idx: HashMap<Cow<'static, str>, usize>,
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
// pub valid_indices: HashSet<usize>,
}
impl Charset {
// 内部底层统一收拢构造
pub fn new(tokens: Vec<Cow<'static, str>>) -> Self {
let mut char_to_idx = HashMap::with_capacity(tokens.len());
for (idx, token) in tokens.iter().enumerate() {
char_to_idx.entry(token.clone()).or_insert(idx);
// 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑)
// char_to_idx.entry(token.to_string()).or_insert(idx);
}
Self {
tokens,
char_to_idx,
}
}
// --- 业务策略方法 ---
/// 将字符转为索引，不存在返回 -1 (保持与原 Python 库行为一致)
pub fn char_to_index(&self, char_str: &str) -> i32 {
if let Some(&idx) = self.char_to_idx.get(char_str) {
idx as i32
} else {
-1
}
}
/// 将索引转为字符引用，零拷贝。若越界返回 None
pub fn index_to_char_ref(&self, index: usize) -> Option<&str> {
self.tokens.get(index).map(|cow| cow.as_ref())
}
pub fn is_valid_char(&self, char_str: &str) -> bool {
self.char_to_idx.get(char_str).is_some()
}
pub fn size(&self) -> usize {
self.tokens.len()
}
}
// ==========================================
// 4. 标准 Display 接口实现 (对应 __str__)
// ==========================================
impl std::fmt::Display for Charset {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Charset [Total Size: {}", self.size(),)
}
}
// /// 从外部外部 JSON 文件动态加载字符集(在后续优化中移除)
// pub fn from_json_file<P: AsRef<Path>>(path: P) -> anyhow::Result<Self> {
// let path = path.as_ref();
// if !path.exists() {
// return Err(anyhow!("模型元数据配置文件不存在: {:?}", path));
// }
//
// let mut file = File::open(path)?;
// let mut content = String::new();
// file.read_to_string(&mut content)?;
//
// let dto: ModelMetadataDto = serde_json::from_str(&content)
// .map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?;
//
// // 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,
// })
// }

38
tests/ocr_test.rs
View File

@@ -1,8 +1,11 @@
use ddddocr_rs::{Ocr, Slider, Detector, ModelMetadata, OcrSession, DetBuilder, DetSession}; // 假设你的包名是这个
use ddddocr_rs::models::det::DetectionResult;
use ddddocr_rs::{DetBuilder, DetSession, Detector, ModelMetadata, Ocr, OcrSession, Slider}; // 假设你的包名是这个
use image::{DynamicImage, Rgb};
use std::fs;
use std::path::Path;
use ddddocr_rs::models::det::DetectionResult;
mod char_slice;
use char_slice::CHARSET_BETA;
use ddddocr_rs::models::ocr::metadata::{Normalization, Resize};
fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
// 1. 先将泛型转为具体的 &Path 引用
@@ -16,8 +19,8 @@ fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
}
/// 将检测结果绘制在图像上并保存
fn save_debug_image(
dynamic_img: &DynamicImage, // 【优化点 1】直接传入解码好的引用，拒绝重复解码
bboxes: &[DetectionResult], // 【修改点 1】类型改为自定义结构体切片
dynamic_img: &DynamicImage, // 【优化点 1】直接传入解码好的引用，拒绝重复解码
bboxes: &[DetectionResult], // 【修改点 1】类型改为自定义结构体切片
output_path: &str,
) -> anyhow::Result<()> {
// 删除了原本的 let dynamic_img = image::load_from_memory(image_bytes)?;
@@ -60,17 +63,29 @@ fn save_debug_image(
Ok(())
}
#[test]
fn test_full_classification() {
// 1. 初始化模型
let ocr = OcrSession::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_static_slice(
CHARSET_BETA,
false,
Resize::DynamicWidth(64),
1,
Normalization::MinusOneToOne,
),
)
.expect("模型加载失败");
// 2. 加载测试图片
let img = image::open("samples/code2.png").expect("测试图片不存在");
// 3. 执行识别
let result = Ocr::new(&ocr).predict(&img).expect("识别过程出错").into_text();
let result = Ocr::new(&ocr)
.predict(&img)
.expect("识别过程出错")
.into_text();
println!("识别结果: {}", result);
assert!(!result.is_empty());
@@ -101,10 +116,7 @@ fn test_det_load() -> anyhow::Result<()> {
for (i, bbox) in bboxes.iter().enumerate() {
// 【修改点 3】将原来的 bbox[0].. 索引访问改为结构体字段访问
println!(
"目标 [{}]: x1={}, y1={}, x2={}, y2={}, 分数={:.4}, 类别ID={}",
i, bbox.x1, bbox.y1, bbox.x2, bbox.y2, bbox.score, bbox.class_id
);
println!("目标 [{}]: {}", i, bbox);
}
}
Ok(())
@@ -129,9 +141,7 @@ fn test_real_slide_match() {
// 3. 打印结果
println!("-------------------------------------------");
println!("滑块匹配测试结果:");
println!("检测坐标: [x: {}, y: {}]", result.target_x, result.target_y);
println!("置信度: {:.4}", result.confidence);
println!("{}", result);
println!("耗时: {:?}", duration);
println!("-------------------------------------------");