diff --git a/src/charset.rs b/src/charset.rs index 9cee79b..52a5805 100644 --- a/src/charset.rs +++ b/src/charset.rs @@ -544,7 +544,8 @@ pub enum CharsetRestrict { // Single(String), /// 直接设置完整的 Token 白名单（支持多字 Token），例如 vec!["html".to_string()] CustomList(Vec), - + /// 完美对应 Python 传入 int 时的行为：截取并只保留模型字符集前 N 个字符 + TopN(usize), /// 核心组合子：满足左边或右边任意一个条件即可（即 A + B 的并集逻辑） /// 使用 Box 打破 Rust 编译期对递归枚举的无限大小限制 Or(Box, Box), @@ -577,6 +578,7 @@ impl CharsetRestrict { 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::TopN(_) => false, CharsetRestrict::Or(left, right) => left.matches(s) || right.matches(s), } } @@ -638,7 +640,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 is_valid_char(&self, char_str: &str) -> bool { self.char_to_idx.get(char_str).is_some() @@ -646,7 +648,7 @@ impl Charset { pub fn size(&self) -> usize { self.tokens.len() } - + } // ========================================== diff --git a/src/lib.rs b/src/lib.rs index 4c89e61..29e3d8c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,7 +9,7 @@ use image::DynamicImage; use std::fmt::{Display, Formatter}; // 关键点：直接使用 tract 重导出的 ndarray -use crate::charset::get_default_charset; +use crate::charset::{get_default_charset, CharsetRestrict}; use crate::models::ocr::ColorRange; use models::det::Det; use models::loader::ModelSession; @@ -96,7 +96,8 @@ impl Display for DdddOcr { impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { - Runtime::Ocr(s) => s.predict(img).run(), + // Runtime::Ocr(s) => s.predict(img).run(), + Runtime::Ocr(s) => s.builder().charset_restrict(&CharsetRestrict::Digit).predict(img), Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), } } diff --git a/src/models/ocr.rs b/src/models/ocr.rs index b8e3d16..6071fe5 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -35,105 +35,11 @@ impl Ocr { model_metadata, }) } - pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> { - OcrBuilder::new(self, image) - } -} - -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_restrict: Option, -} - -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_restrict: None, - } - } - pub fn png_fix(mut self, value: bool) -> Self { - self.png_fix = value; - self - } - pub fn color_filter_colors(mut self, value: Vec) -> Self { - self.color_filter_colors = Some(value); - self - } - pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { - self.color_filter_custom_ranges = Some(value); - self - } - pub fn charset_restrict(mut self, restrict: CharsetRestrict) -> Self { - self.charset_restrict = Some(restrict); - self - } - - pub fn run(&self) -> anyhow::Result { - let tensor = self.preprocess_image(self.image, self.png_fix)?; - - 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 - /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 - fn preprocess_image(&self, img: &DynamicImage, png_fix: bool) -> anyhow::Result { - // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) - let _ = if png_fix && img.color().has_alpha() { - png_rgba_white_preprocess(img) - } else { - img.clone() - }; - - let h = 64u32; - let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32; - let gray_img = convert_to_grayscale(img); - let resized = resize_image(&gray_img, w, h); - // resized.save("debug_preprocessed.png").unwrap(); - // 1. 预处理：转灰度 -> Resize -> 归一化 - // let resized = img.resize_exact(w, h, FilterType::Lanczos3).to_luma8(); - - // 使用 tract_ndarray 构造，避免版本冲突 - let array = - tract_ndarray::Array4::from_shape_fn((1, 1, h as usize, w as usize), |(_, _, y, x)| { - let pixel = resized.get_pixel(x as u32, y as u32)[0] as f32; - (pixel / 255.0 - 0.5) / 0.5 - }); - - let tensor = Tensor::from(array); - - Ok(tensor) - } /// 对应 Python 的 _inference fn inference(&self, tensor: Tensor) -> anyhow::Result { // tract 的 run 会返回一个 Vec，我们通常只需要第一个输出 // let result = self.session.run(tvec!(tensor.into()))?; let mut result = self - .ocr .session .run(tvec!(tensor.into())) .context("执行模型推理失败")?; @@ -205,70 +111,189 @@ impl<'a> OcrBuilder<'a> { )), } } - /// 获取有效字符索引列表 (用于外部验证或过滤) - pub fn get_valid_indices(&self) -> HashSet { - let (_, valid_indices) = self.valid_indices(); - valid_indices - } - fn valid_indices(&self) -> (bool, HashSet) { + pub fn builder(&'_ self) -> OcrBuilder<'_> { + OcrBuilder::new(self) + } +} + +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_restrict: Option>, +} + +impl<'a> OcrBuilder<'a> { + // 初始化任务，设置默认参数 + pub fn new(ocr: &'a Ocr) -> Self { + Self { + ocr, + // image, + png_fix: false, // 默认值 + probability: false, + color_filter_colors: None, + color_filter_custom_ranges: None, + charset_restrict: None, + } + } + pub fn png_fix(mut self, value: bool) -> Self { + self.png_fix = value; + self + } + pub fn color_filter_colors(mut self, value: Vec) -> Self { + self.color_filter_colors = Some(value); + self + } + pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { + self.color_filter_custom_ranges = Some(value); + self + } + pub fn charset_restrict(mut self, restrict: &CharsetRestrict) -> Self { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; - /// 根据传入的 CharsetRestrict 枚举策略，动态更新有效索引 - // 1. 🧠 零克隆防御战：在局部判断并动态构建专属的白名单判定表 - let mut valid_indices = HashSet::new(); + // let mut temp_indices = Vec::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); + let estimated_capacity = match restrict { + CharsetRestrict::Digit => 16, + CharsetRestrict::Lowercase | CharsetRestrict::Uppercase => 32, + CharsetRestrict::CustomList(vec) => vec.len() + 1, // 动态匹配列表大小 + CharsetRestrict::TopN(n) => *n + 1, + _ => 128, // 组合子（Or）等复杂情况，给个 128 黄金保底值 + }; + // 精准开辟内存，完美避开 8210 个槽位的巨大空置浪费 + let mut temp_indices = Vec::with_capacity(estimated_capacity); + + if let CharsetRestrict::TopN(n) = *restrict { + let limit = std::cmp::min(n, tokens.len()); + // 边界防御：CTC Blank (索引 0) 必须无条件放行 + temp_indices.push(0); + // temp_indices.extend(0..limit); + // has_any_match = limit > &0; + // 塞入剩余的有效索引范围（排除0，从1开始截取） + if limit > 1 { + temp_indices.extend(1..limit); + has_any_match = true; + } + } else { 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); + temp_indices.push(idx); + } else if restrict.matches(token_str) { + temp_indices.push(idx); has_any_match = true; } } - - // 终极防御：如果除了 Blank 之外，没有任何一个字符被匹配到 - if !has_any_match { - valid_indices = (0..tokens.len()).collect(); - println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); - } } - (has_any_match, valid_indices) + // self.charset_restrict = Some(restrict); + // 终极防御：如果除了 Blank 外什么都没匹配上，退化恢复为 None（全量识别） + if !has_any_match { + println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); + self.charset_restrict = None; + } else { + // 这一步非常重要：二分查找（binary_search）强依赖数组【有序】。 + // TopN 天然有序，但如果是用户自定义的 CustomList 或者复杂的 Or 组合， + // 遍历出来的索引天然有序，但为了绝对的安全，我们在这里顺手排个序 + temp_indices.sort_unstable(); + self.charset_restrict = Some(temp_indices); + } + self } +} +impl<'a> OcrBuilder<'a> { + pub fn predict(&self, image: &DynamicImage) -> anyhow::Result { + let tensor = self.preprocess_image(image)?; + + let raw_tensor = self.ocr.inference(tensor)?; + let raw_indices = self.ocr.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 + /// 负责：透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 + fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result { + // A. 修复 PNG 透明背景 (内部逻辑你之前已实现) + let _ = if self.png_fix && img.color().has_alpha() { + png_rgba_white_preprocess(img) + } else { + img.clone() + }; + + let h = 64u32; + let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32; + let gray_img = convert_to_grayscale(img); + let resized = resize_image(&gray_img, w, h); + // resized.save("debug_preprocessed.png").unwrap(); + // 1. 预处理：转灰度 -> Resize -> 归一化 + // let resized = img.resize_exact(w, h, FilterType::Lanczos3).to_luma8(); + + // 使用 tract_ndarray 构造，避免版本冲突 + let array = + tract_ndarray::Array4::from_shape_fn((1, 1, h as usize, w as usize), |(_, _, y, x)| { + let pixel = resized.get_pixel(x as u32, y as u32)[0] as f32; + (pixel / 255.0 - 0.5) / 0.5 + }); + + let tensor = Tensor::from(array); + + Ok(tensor) + } +} +impl<'a> OcrBuilder<'a> { + pub fn get_valid_indices(&self) -> HashSet { + match &self.charset_restrict { + Some(indices) => indices.iter().cloned().collect(), + // 如果是 None，现场映射出全量索引集给外部 + None => (0..self.ocr.model_metadata.charset.tokens.len()).collect(), + } + } + // compute_valid_indices + // fn valid_indices(&self) -> (bool, HashSet) { + // let charset = &self.ocr.model_metadata.charset; + /// 🌟 【按需延迟打印】：当用户真的需要“知道当前有哪些限制字符”时，一秒反查并打印 /// 这里的 &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() + match &self.charset_restrict { + Some(indices) => indices + .iter() + .filter_map(|&idx| tokens.get(idx).map(|cow| cow.as_ref())) + .collect(), + // 如果是 None，现场映射出全量 Token 视图给外部 + None => tokens.iter().map(|cow| cow.as_ref()).collect(), + } } pub fn valid_size(&self) -> usize { - self.get_valid_indices().len() + match &self.charset_restrict { + Some(indices) => indices.len(), + None => self.ocr.model_metadata.charset.tokens.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; @@ -291,11 +316,11 @@ impl<'a> OcrBuilder<'a> { Err(_) => continue, }; - // 4. 终极性能：既然你的有效索引库必然有全量或部分数据，这里直接进行 O(1) 包含校验 - // 注意：去掉原本错误的 `!` - if has_any_match { - if !valid_indices.contains(&u_idx) { - continue; // 不在有效字符集内，安全跳过 + // 史诗级加速点：如果是 None，说明没限制，根本不进入分支，直接放行！ + // 只有当有具体限制（Some）时，才去跑 4-5 次 CPU 寄存器级别的二分查找 + if let Some(ref indices) = self.charset_restrict { + if indices.binary_search(&u_idx).is_err() { + continue; } }