diff --git a/src/charset.rs b/src/charset.rs index c5242d5..9cee79b 100644 --- a/src/charset.rs +++ b/src/charset.rs @@ -600,11 +600,11 @@ impl Add for CharsetRestrict { #[derive(Debug, Clone)] pub struct Charset { // 使用 Cow 统一静态切片和动态读取的 Vec，内部实现真正的零拷贝 - tokens: Vec>, + pub tokens: Vec>, // 反向查找表，保证字符转索引为 O(1) - char_to_idx: HashMap, usize>, + pub char_to_idx: HashMap, usize>, // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速) - valid_indices: HashSet, + // pub valid_indices: HashSet, } impl Charset { @@ -617,48 +617,14 @@ impl Charset { // 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) { @@ -672,21 +638,7 @@ impl Charset { 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() @@ -694,10 +646,7 @@ impl Charset { pub fn size(&self) -> usize { self.tokens.len() } - - pub fn valid_size(&self) -> usize { - self.valid_indices.len() - } + } // ========================================== @@ -707,9 +656,8 @@ 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: {}]", + "Charset [Total Size: {}", self.size(), - self.valid_size() ) } } diff --git a/src/lib.rs b/src/lib.rs index 3965895..4c89e61 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,6 +14,7 @@ use crate::models::ocr::ColorRange; use models::det::Det; use models::loader::ModelSession; use models::ocr::Ocr; +use crate::model_metadata::ModelMetadata; pub enum ModelSpec { /// 默认 OCR (使用内置路径) @@ -22,7 +23,7 @@ pub enum ModelSpec { /// 自定义 OCR (路径由用户提供) CustomOcrModel { path: String, - charset: Vec, + model_metadata: ModelMetadata, }, } impl ModelSpec { @@ -61,9 +62,9 @@ impl DdddOcrBuilder { } /// 设置自定义 OCR 路径 - pub fn custom_ocr(mut self, path: String, charset: Vec) -> Self { + pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self { // 直接重写枚举，替换掉之前的 Ocr 或 Det - self.mode = ModelSpec::CustomOcrModel { path, charset }; + self.mode = ModelSpec::CustomOcrModel { path, model_metadata }; self } @@ -72,10 +73,10 @@ impl DdddOcrBuilder { let runtime = match self.mode { ModelSpec::OcrModel => Runtime::Ocr(Ocr::new( ModelSpec::DEFAULT_OCR_PATH.into(), - get_default_charset(), + ModelMetadata::from_builtin_beta(), )?), ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?), - ModelSpec::CustomOcrModel { path, charset } => Runtime::Ocr(Ocr::new(path, charset)?), + ModelSpec::CustomOcrModel { path, model_metadata } => Runtime::Ocr(Ocr::new(path, model_metadata)?), }; Ok(DdddOcr { runtime }) diff --git a/src/model_metadata.rs b/src/model_metadata.rs index efd7f8d..24bd590 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -34,7 +34,7 @@ struct ModelMetadataDto { #[derive(Debug, Clone)] pub struct ModelMetadata { /// 字符集管理器 - pub charset: Charset, + pub charset: Charset, /// 是否为单字识别模型 pub word: bool, /// 预处理的缩放策略 diff --git a/src/models/ocr.rs b/src/models/ocr.rs index aa7266b..b8e3d16 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -1,23 +1,23 @@ +use crate::charset::CharsetRestrict; +use crate::model_metadata::ModelMetadata; use crate::models::base::ModelArgs; use crate::models::loader::{ModelLoader, ModelSession, ModelType}; use crate::utils::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use anyhow::Context; use image::DynamicImage; +use std::collections::HashSet; 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)); - - pub struct Ocr { pub session: RunnableModel, Graph>>, - pub charset: Vec, + pub model_metadata: ModelMetadata, } impl ModelSession for Ocr { fn get_model_type(&self) -> ModelType { @@ -28,9 +28,12 @@ impl ModelSession for Ocr { } } impl Ocr { - pub fn new(model_path: String, charset: Vec) -> Result { + pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result { let session = ModelLoader::load_model(&model_path)?.session; - Ok(Self { session, charset }) + Ok(Self { + session, + model_metadata, + }) } pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> { OcrBuilder::new(self, image) @@ -50,7 +53,7 @@ pub struct OcrBuilder<'a> { /// 颜色过滤：自定义RGB范围 color_filter_custom_ranges: Option>, /// 字符集范围 - charset_range: Option, + charset_restrict: Option, } impl<'a> OcrBuilder<'a> { @@ -63,7 +66,7 @@ impl<'a> OcrBuilder<'a> { probability: false, color_filter_colors: None, color_filter_custom_ranges: None, - charset_range: None + charset_restrict: None, } } pub fn png_fix(mut self, value: bool) -> Self { @@ -78,21 +81,22 @@ impl<'a> OcrBuilder<'a> { self.color_filter_custom_ranges = Some(value); self } - pub fn charset_range(mut self, range: CharsetRange) -> Self { - self.charset_range = Some(range); + pub fn charset_restrict(mut self, restrict: CharsetRestrict) -> Self { + self.charset_restrict = Some(restrict); self } - pub fn run(&self) -> Result { + pub fn run(&self) -> anyhow::Result { let tensor = self.preprocess_image(self.image, self.png_fix)?; - // - // let result = self.session.run(tvec!(tensor.into()))?; - // // 3. 解析结果 - // // let output = result[0].to_array_view::()?; - let output = self.inference(tensor)?; - let output2 = self.process_text_output(&output)?; - Ok(self.ctc_decode_indices(&output2)) - // Ok("ocr result".to_string()) + + let raw_tensor = self.inference(tensor)?; + let raw_indices = self.extract_indices_from_tensor(&raw_tensor)?; + + // 步骤 2: 将索引切片 `&[i64]` 传给解码器进行 CTC 去重和字符映射 + let final_text = self.ctc_decode_to_string(&raw_indices); + + println!("最终识别出的验证码是: {}", final_text); + Ok(final_text) } /// 对应 Python 的 _preprocess_image @@ -136,8 +140,9 @@ impl<'a> OcrBuilder<'a> { println!("模型输出原始数据: {:?}", result); Ok(result.remove(0).into_tensor()) } + /// 核心解析逻辑：将模型输出的各种维度/类型的 Tensor 转为字符索引序列 - fn process_text_output(&self, raw_tensor: &Tensor) -> anyhow::Result> { + fn extract_indices_from_tensor(&self, raw_tensor: &Tensor) -> anyhow::Result> { let shape = raw_tensor.shape(); println!("模型输出shape数据: {:?}", shape); let datum_type = raw_tensor.datum_type(); @@ -172,6 +177,9 @@ impl<'a> OcrBuilder<'a> { // 形状: [Steps, Classes] -> 已经剥离了 Batch 维度 (shape[0], shape[1], view.into_dyn()) } + // 形状: [Classes] -> 单字符输出（对应 Python 的 ndim == 0 保护逻辑） + // 我们把它虚构成一个 [1, Classes] 的 2D 矩阵来复用后面的 argmax 逻辑 + 1 => (1, shape[0], view.into_dyn()), _ => return Err(anyhow::anyhow!("不支持的输出维度: {:?}", shape)), }; let array_2d = data_view.to_shape((steps, classes))?; @@ -197,29 +205,107 @@ impl<'a> OcrBuilder<'a> { )), } } - fn ctc_decode_indices(&self, predicted_indices: &[i64]) -> String { + /// 获取有效字符索引列表 (用于外部验证或过滤) + pub fn get_valid_indices(&self) -> HashSet { + let (_, valid_indices) = self.valid_indices(); + valid_indices + } + + fn valid_indices(&self) -> (bool, HashSet) { + let charset = &self.ocr.model_metadata.charset; + let tokens = &charset.tokens; + /// 根据传入的 CharsetRestrict 枚举策略，动态更新有效索引 + // 1. 🧠 零克隆防御战：在局部判断并动态构建专属的白名单判定表 + let mut valid_indices = HashSet::new(); + let mut has_any_match = false; + if let Some(ref policy) = self.charset_restrict { + // 🧠 性能精算：根据限制策略的类型，智能分配合适的初始容量，1个字节都不浪费！ + let estimated_capacity = match policy { + CharsetRestrict::Digit => 16, + CharsetRestrict::Lowercase | CharsetRestrict::Uppercase => 32, + CharsetRestrict::CustomList(vec) => vec.len() + 1, // 动态匹配列表大小 + _ => 128, // 组合子（Or）等复杂情况，给个 128 黄金保底值 + }; + // 🚀 精准开辟内存，完美避开 8120 个槽位的巨大空置浪费 + valid_indices = HashSet::with_capacity(estimated_capacity); + + for (idx, token) in tokens.iter().enumerate() { + let token_str = token.as_ref(); + // CTC Blank 空字符串无条件放行，其余交给超高性能的 matches + if token_str.is_empty() { + valid_indices.insert(idx); + } else if policy.matches(token_str) { + valid_indices.insert(idx); + has_any_match = true; + } + } + + // 终极防御：如果除了 Blank 之外，没有任何一个字符被匹配到 + if !has_any_match { + valid_indices = (0..tokens.len()).collect(); + println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); + } + } + (has_any_match, valid_indices) + } + /// 🌟 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 + /// 这里的 &str 完美借用了自 tokens，依然是彻底的零拷贝！ + pub fn get_valid_tokens(&self) -> Vec<&str> { + let charset = &self.ocr.model_metadata.charset; + let tokens = &charset.tokens; + self.get_valid_indices() + .iter() + .map(|&idx| tokens[idx].as_ref()) + .collect() + } + pub fn valid_size(&self) -> usize { + self.get_valid_indices().len() + } + fn ctc_decode_to_string(&self, predicted_indices: &[i64]) -> String { println!("indices模型输出原始数据: {:?}", predicted_indices); + let charset = &self.ocr.model_metadata.charset; + let tokens = &charset.tokens; + // let valid_indices = &charset.valid_indices; + + let (has_any_match, valid_indices) = self.valid_indices(); // 对应 _ctc_decode_indices 的逻辑：去重、去 blank (0) let mut res = String::new(); let mut prev_idx: i64 = -1; for &idx in predicted_indices { - // 1. 跳过连续重复的索引 - // 2. 跳过 blank 字符 (假设索引 0 是 blank) - if idx != prev_idx && idx != 0 { - if let Ok(u_idx) = usize::try_from(idx) { - if let Some(char_str) = self.ocr.charset.get(u_idx) { - res.push_str(char_str); - } else { - // 保护逻辑：如果模型预测的索引超出了字符集范围 - eprintln!("警告: 预测索引 {} 超出字符集范围", u_idx); - } + // 1. CTC 去重：如果是连续重复的，直接跳过 + if idx == prev_idx { + continue; + } + // 【关键核心】只要不是连续重复，立刻更新 prev_idx 状态，绝对不能被后续的过滤短路！ + prev_idx = idx; + + // 2. CTC 过滤 Blank (0) + if idx == 0 { + continue; + } + // 3. 类型安全转换 + let u_idx = match usize::try_from(idx) { + Ok(u) => u, + Err(_) => continue, + }; + + // 4. 终极性能：既然你的有效索引库必然有全量或部分数据，这里直接进行 O(1) 包含校验 + // 注意：去掉原本错误的 `!` + if has_any_match { + if !valid_indices.contains(&u_idx) { + continue; // 不在有效字符集内，安全跳过 } } - prev_idx = idx; + + // 5. 字符映射 + if let Some(char_str) = tokens.get(u_idx) { + res.push_str(char_str); + } else { + eprintln!("警告: 预测索引 {} 超出字符集范围", u_idx); + } } - println!("最终识别出的验证码是: {}", res); res } } diff --git a/tests/ocr_test.rs b/tests/ocr_test.rs index 928eee6..d6db696 100644 --- a/tests/ocr_test.rs +++ b/tests/ocr_test.rs @@ -66,7 +66,7 @@ fn test_full_classification() { let ocr = DdddOcrBuilder::new().build().expect("模型加载失败"); // 2. 加载测试图片 - let img = image::open("samples/code2.png").expect("测试图片不存在"); + let img = image::open("samples/code3.png").expect("测试图片不存在"); // 3. 执行识别 let result = ocr.classification(&img).expect("识别过程出错");