diff --git a/Cargo.toml b/Cargo.toml index 24fac8c..55e6c95 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,4 +9,6 @@ tract-onnx = { version = "0.21.10" } anyhow = "1.0.102" image = "0.25.10" base64 = "0.22.1" -imageproc = { version = "0.26.2", default-features = true } \ No newline at end of file +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 diff --git a/src/charset.rs b/src/charset.rs index fc44cef..c5242d5 100644 --- a/src/charset.rs +++ b/src/charset.rs @@ -514,6 +514,202 @@ pub const CHARSET_BETA: &[&str] = &[ "谬", "溝", "言", "哽", "婿", "猿", "跗", "獴", "俜", "呙", "弗", "凿", "窭", "铌", "友", "唉", "怫", "荘", ]; + +pub const CHARSET_OLD: &[&str] = &["", "笤", "谴", "膀", "荔"]; + pub fn get_default_charset() -> Vec { CHARSET_BETA.iter().map(|&s| s.to_string()).collect() -} \ No newline at end of file +} + +use std::borrow::Cow; +use std::collections::{HashMap, HashSet}; + +use std::ops::{Add, Deref}; +// 字符集范围类型 +/// 字符集范围限制组合子枚举 +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CharsetRestrict { + /// 纯整数 0-9 + Digit, + + /// 纯小写字母 a-z + Lowercase, + + /// 纯大写字母 A-Z + Uppercase, + + // /// 过滤模式：删除所有 ASCII 字母和数字（通常用于仅保留汉字、特殊标点） + // // ExcludeAlphanumeric, + // /// 自定义单字字符集，例如 "0123456789+-x/=" + // Single(String), + /// 直接设置完整的 Token 白名单（支持多字 Token），例如 vec!["html".to_string()] + CustomList(Vec), + + /// 核心组合子：满足左边或右边任意一个条件即可（即 A + B 的并集逻辑） + /// 使用 Box 打破 Rust 编译期对递归枚举的无限大小限制 + Or(Box, Box), +} +impl From for CharsetRestrict { + fn from(value: i32) -> Self { + match value { + 0 => Self::Digit, + 1 => Self::Lowercase, + 2 => Self::Uppercase, + // 3 => Self::LowercaseUppercase, + // 4 => Self::LowercaseDigit, + // 5 => Self::UppercaseDigit, + // 6 => Self::LowercaseUppercaseDigit, + // 7 => Self::DefaultCharsetLowercaseUppercaseDigit, + _ => panic!("invalid charset range: {}", value), + } + } +} +impl CharsetRestrict { + /// 💡 辅助构造函数：直接在源头把用户的长字符串切碎，伪装成基础积木 + pub fn from_chars(custom_str: &str) -> Self { + let tokens = custom_str.chars().map(|c| c.to_string()).collect(); + CharsetRestrict::CustomList(tokens) + } + // 内部递归收集器：利用硬编码切片快速无损展开 + pub(crate) fn matches(&self, s: &str) -> bool { + match self { + CharsetRestrict::Digit => s.len() == 1 && s.as_bytes()[0].is_ascii_digit(), + CharsetRestrict::Lowercase => s.len() == 1 && s.as_bytes()[0].is_ascii_lowercase(), + CharsetRestrict::Uppercase => s.len() == 1 && s.as_bytes()[0].is_ascii_uppercase(), + CharsetRestrict::CustomList(vec) => vec.iter().any(|t| t == s), + CharsetRestrict::Or(left, right) => left.matches(s) || right.matches(s), + } + } +} +// ===================================================================== +// 5. 优雅的魔法：重载 + 运算符 (实现 std::ops::Add) +// ===================================================================== + +/// 支持 `CharsetRestrict::Digit + CharsetRestrict::Lowercase` +impl Add for CharsetRestrict { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + CharsetRestrict::Or(Box::new(self), Box::new(rhs)) + } +} + +// ========================================== +// 3. 字符集核心结构体 (重命名为 Charset) +// ========================================== +#[derive(Debug, Clone)] +pub struct Charset { + // 使用 Cow 统一静态切片和动态读取的 Vec，内部实现真正的零拷贝 + tokens: Vec>, + // 反向查找表，保证字符转索引为 O(1) + char_to_idx: HashMap, usize>, + // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速) + valid_indices: HashSet, +} + +impl Charset { + // 内部底层统一收拢构造 + pub fn new(tokens: Vec>) -> 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); + } + + // 默认初始化时，所有索引均为有效状态 + let valid_indices = (0..tokens.len()).collect(); + + Self { + tokens, + char_to_idx, + valid_indices, + } + } + + // --- 业务策略方法 --- + + /// 根据传入的 CharsetRange 枚举策略，动态更新有效索引 + pub fn apply_range_policy(&mut self, policy: &CharsetRestrict) -> bool { + let mut has_any_match = false; + // 3. 清空原有的索引 + self.valid_indices.clear(); + // 4. 执行 O(1) 级别的求交集过滤 + for (idx, token) in self.tokens.iter().enumerate() { + let token_str = token.as_ref(); + // CTC Blank 空字符串无条件放行，其余交给超高性能的 matches + if token_str.is_empty() { + self.valid_indices.insert(idx); + } else if policy.matches(token_str) { + self.valid_indices.insert(idx); + has_any_match = true; + } + } + // 🛡️ 终极防御：如果除了 Blank 之外，没有任何一个字符被匹配到 + if !has_any_match { + // 策略 C：智能降级，一键恢复全量字符集，防止模型“交白卷” + self.reset_range_policy(); + println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); + return false; // 返回 false 提示外部：策略未实际生效，已降级 + } + true + } + /// 清除范围限制，恢复完整字符集 + pub fn reset_range_policy(&mut self) { + self.valid_indices = (0..self.tokens.len()).collect(); + } + + /// 将字符转为索引，不存在返回 -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 get_valid_indices(&self) -> Vec { + let mut indices: Vec = self.valid_indices.iter().copied().collect(); + indices.sort_unstable(); + indices + } + /// 🌟 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 + /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ + pub fn get_valid_tokens(&self) -> Vec<&str> { + self.get_valid_indices() + .iter() + .map(|&idx| self.tokens[idx].as_ref()) + .collect() + } + + 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() + } + + pub fn valid_size(&self) -> usize { + self.valid_indices.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: {}, Active Range Size: {}]", + self.size(), + self.valid_size() + ) + } +} diff --git a/src/lib.rs b/src/lib.rs index 8899498..3965895 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,6 +2,7 @@ mod charset; pub mod models; pub mod utils; +mod model_metadata; use anyhow::Result; use image::DynamicImage; @@ -94,7 +95,7 @@ impl Display for DdddOcr { impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { - Runtime::Ocr(s) => s.predict(img, false), + Runtime::Ocr(s) => s.predict(img).run(), Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), } } diff --git a/src/model_metadata.rs b/src/model_metadata.rs new file mode 100644 index 0000000..efd7f8d --- /dev/null +++ b/src/model_metadata.rs @@ -0,0 +1,122 @@ +use crate::charset::{CHARSET_BETA, CHARSET_OLD, Charset, CharsetRestrict}; +use anyhow::{Result, anyhow}; +use serde::Deserialize; +use std::borrow::Cow; +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::Read; +use std::path::Path; +// ===================================================================== +// 1. 辅助定义的枚举与结构体 +// ===================================================================== + +/// 图像缩放策略枚举 +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Resize { + /// 固定宽高，例如 (64, 64) + Fixed(u32, u32), + /// 高度固定，宽度根据原始比例动态计算（对应 Python 的 [-1, H]） + DynamicWidth(u32), + /// 单字识别的正方形切图（对应 Python 的 word 为 True 且 [-1, H]） + Square(u32), +} + +/// 仅用于反序列化 JSON 的中间临时结构体（DTO） +#[derive(Deserialize)] +struct ModelMetadataDto { + charset: Vec, + word: bool, + #[serde(alias = "image")] + resize: Vec, + channel: u8, +} + +#[derive(Debug, Clone)] +pub struct ModelMetadata { + /// 字符集管理器 + pub charset: Charset, + /// 是否为单字识别模型 + pub word: bool, + /// 预处理的缩放策略 + pub resize: Resize, + /// 图像通道数 (1 或 3) + pub channel: u8, +} + +impl ModelMetadata { + // --- 优雅的工厂模式构造器 --- + /// 从预设的旧版字符集创建 + pub fn from_builtin_old() -> Self { + Self::from_static_slice(CHARSET_OLD, false, Resize::Fixed(64, 64), 1) + } + + /// 从预设的 Beta 版字符集创建 + pub fn from_builtin_beta() -> Self { + Self::from_static_slice(CHARSET_BETA, false, Resize::Fixed(64, 64), 1) + } + + /// 通用的静态切片转换构造器 + pub fn from_static_slice( + slice: &[&'static str], + word: bool, + resize: Resize, + channel: u8, + ) -> Self { + let tokens: Vec> = slice.iter().map(|&s| Cow::Borrowed(s)).collect(); + Self { + charset: Charset::new(tokens), + word, + resize, + channel, + } + } + + /// 从外部外部 JSON 文件动态加载字符集 + pub fn from_json_file>(path: P) -> Result { + 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> = + 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, + }) + } +} diff --git a/src/models/ocr.rs b/src/models/ocr.rs index e287b5f..aa7266b 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -8,92 +8,13 @@ use tract_onnx::prelude::tract_ndarray::s; use tract_onnx::prelude::{ DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec, }; +use crate::charset::CharsetRange; // 颜色过滤的自定义范围：(低值RGB, 高值RGB) pub type ColorRange = ((u8, u8, u8), (u8, u8, u8)); -// 字符集范围类型 -#[derive(Debug, Clone)] -pub enum CharsetRange { - All, // 所有字符 - Digit, // 数字 - Letter, // 字母 - Alphanumeric, // 字母数字 - Single(String), // 单字符串 - Multiple(Vec), // 多个字符串 - Range(char, char), // 字符范围 - Custom(Vec), // 自定义字符列表 -} -#[derive(Debug, Clone)] -pub struct PredictArgs { - /// 是否修复PNG格式问题 - pub png_fix: bool, - /// 是否返回概率信息 - pub probability: bool, - /// 颜色过滤：保留的颜色列表 - pub color_filter_colors: Option>, - /// 颜色过滤：自定义RGB范围 - pub color_filter_custom_ranges: Option>, - /// 字符集范围 - pub charset_range: Option, -} -impl Default for PredictArgs { - fn default() -> Self { - Self { - png_fix: false, - probability: false, - color_filter_colors: None, - color_filter_custom_ranges: None, - charset_range: None, - } - } -} -impl PredictArgs { - pub fn new() -> Self { - Self::default() - } - - // Builder 模式方法 - pub fn png_fix(mut self, enabled: bool) -> Self { - self.png_fix = enabled; - self - } - - pub fn probability(mut self, enabled: bool) -> Self { - self.probability = enabled; - self - } - - pub fn color_filter_colors(mut self, colors: Vec) -> Self { - self.color_filter_colors = Some(colors); - self - } - - pub fn color_filter_custom_ranges(mut self, ranges: Vec) -> Self { - self.color_filter_custom_ranges = Some(ranges); - self - } - - pub fn charset_range(mut self, range: CharsetRange) -> Self { - self.charset_range = Some(range); - self - } - - // 便捷构造方法 - pub fn quick() -> Self { - Self::default() - } - - pub fn with_probability() -> Self { - Self::default().probability(true) - } - - pub fn with_png_fix() -> Self { - Self::default().png_fix(true) - } -} pub struct Ocr { pub session: RunnableModel, Graph>>, pub charset: Vec, @@ -111,28 +32,38 @@ impl Ocr { let session = ModelLoader::load_model(&model_path)?.session; Ok(Self { session, charset }) } - pub fn task<'a>(&'a self, image: &'a DynamicImage) -> OcrTask { - OcrTask::new(self, image) + pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> { + OcrBuilder::new(self, image) } } -pub struct OcrTask<'a> { +pub struct OcrBuilder<'a> { ocr: &'a Ocr, image: &'a DynamicImage, + /// 是否修复PNG格式问题 png_fix: bool, + /// 是否返回概率信息 + #[allow(dead_code)] + probability: bool, + /// 颜色过滤：保留的颜色列表 color_filter_colors: Option>, + /// 颜色过滤：自定义RGB范围 color_filter_custom_ranges: Option>, + /// 字符集范围 + charset_range: Option, } -impl<'a> OcrTask<'a> { +impl<'a> OcrBuilder<'a> { // 初始化任务，设置默认参数 pub fn new(ocr: &'a Ocr, image: &'a DynamicImage) -> Self { Self { ocr, image, png_fix: false, // 默认值 + probability: false, color_filter_colors: None, color_filter_custom_ranges: None, + charset_range: None } } pub fn png_fix(mut self, value: bool) -> Self { @@ -147,9 +78,13 @@ impl<'a> OcrTask<'a> { self.color_filter_custom_ranges = Some(value); self } + pub fn charset_range(mut self, range: CharsetRange) -> Self { + self.charset_range = Some(range); + self + } - pub fn predict(&self, image: &DynamicImage, png_fix: bool) -> Result { - let tensor = self.preprocess_image(image, png_fix)?; + pub fn run(&self) -> Result { + let tensor = self.preprocess_image(self.image, self.png_fix)?; // // let result = self.session.run(tvec!(tensor.into()))?; // // 3. 解析结果