refactor: 优化 OcrBuilder 新增通过索引范围控制有效字符集

- 修改 `charset_restrict`类型修改为Option<Vec<usize>> 并重构同名方法数据处理逻辑
- 优化 `ctc_decode_to_string` 内部复用策略计算，通过 `Option` 结构实现无限制请求的全量免检短路加速
- 新增 `CharsetRestrict`枚举新增变体`TopN(usize)` 实现通过索引范围控制有效字符集

src/charset.rs

@@ -544,7 +544,8 @@ pub enum CharsetRestrict {
     // Single(String),
     /// 直接设置完整的 Token 白名单（支持多字 Token），例如 vec!["html".to_string()]
     CustomList(Vec<String>),
-
+    /// 完美对应 Python 传入 int 时的行为：截取并只保留模型字符集前 N 个字符
+    TopN(usize),
     /// 核心组合子：满足左边或右边任意一个条件即可（即 A + B 的并集逻辑）
     /// 使用 Box 打破 Rust 编译期对递归枚举的无限大小限制
     Or(Box<CharsetRestrict>, Box<CharsetRestrict>),
@@ -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),
         }
     }

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<String> {
         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")),
         }
     }

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<Vec<ColorRange>>,
-    /// 颜色过滤：自定义RGB范围
-    color_filter_custom_ranges: Option<Vec<ColorRange>>,
-    /// 字符集范围
-    charset_restrict: Option<CharsetRestrict>,
-}
-
-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<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 charset_restrict(mut self, restrict: CharsetRestrict) -> Self {
-        self.charset_restrict = Some(restrict);
-        self
-    }
-
-    pub fn run(&self) -> anyhow::Result<String> {
-        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<Tensor> {
-        // 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<Tensor> {
         // tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
         // 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<usize> {
+    pub fn builder(&'_ self) -> OcrBuilder<'_> {
-        let (_, valid_indices) = self.valid_indices();
+        OcrBuilder::new(self)
-        valid_indices
+    }
 }
 
-    fn valid_indices(&self) -> (bool, HashSet<usize>) {
+pub struct OcrBuilder<'a> {
+    ocr: &'a Ocr,
+    // image: &'a DynamicImage,
+    /// 是否修复PNG格式问题
+    png_fix: bool,
+    /// 是否返回概率信息
+    #[allow(dead_code)]
+    probability: bool,
+    /// 颜色过滤：保留的颜色列表
+    color_filter_colors: Option<Vec<ColorRange>>,
+    /// 颜色过滤：自定义RGB范围
+    color_filter_custom_ranges: Option<Vec<ColorRange>>,
+    /// 字符集范围
+    charset_restrict: Option<Vec<usize>>,
+}
+
+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<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 charset_restrict(mut self, restrict: &CharsetRestrict) -> Self {
         let charset = &self.ocr.model_metadata.charset;
         let tokens = &charset.tokens;
-        /// 根据传入的 CharsetRestrict 枚举策略，动态更新有效索引
+        // let mut temp_indices = Vec::new();
-        // 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 restrict {
-            let estimated_capacity = match policy {
             CharsetRestrict::Digit => 16,
             CharsetRestrict::Lowercase | CharsetRestrict::Uppercase => 32,
             CharsetRestrict::CustomList(vec) => vec.len() + 1, // 动态匹配列表大小
+            CharsetRestrict::TopN(n) => *n + 1,
             _ => 128, // 组合子（Or）等复杂情况，给个 128 黄金保底值
         };
-            // 🚀 精准开辟内存，完美避开 8120 个槽位的巨大空置浪费
+        // 精准开辟内存，完美避开 8210 个槽位的巨大空置浪费
-            valid_indices = HashSet::with_capacity(estimated_capacity);
+        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);
+                    temp_indices.push(idx);
-                } else if policy.matches(token_str) {
+                } else if restrict.matches(token_str) {
-                    valid_indices.insert(idx);
+                    temp_indices.push(idx);
                     has_any_match = true;
                 }
             }
-
+        }
-            // 终极防御：如果除了 Blank 之外，没有任何一个字符被匹配到
+        // self.charset_restrict = Some(restrict);
+        // 终极防御：如果除了 Blank 外什么都没匹配上，退化恢复为 None（全量识别）
         if !has_any_match {
-                valid_indices = (0..tokens.len()).collect();
             println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。");
+            self.charset_restrict = None;
+        } else {
+            // 这一步非常重要：二分查找（binary_search）强依赖数组【有序】。
+            // TopN 天然有序，但如果是用户自定义的 CustomList 或者复杂的 Or 组合，
+            // 遍历出来的索引天然有序，但为了绝对的安全，我们在这里顺手排个序
+            temp_indices.sort_unstable();
+            self.charset_restrict = Some(temp_indices);
+        }
+        self
     }
 }
-        (has_any_match, valid_indices)
+impl<'a> OcrBuilder<'a> {
+    pub fn predict(&self, image: &DynamicImage) -> anyhow::Result<String> {
+        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<Tensor> {
+        // 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<usize> {
+        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<usize>) {
+    //     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()
+        match &self.charset_restrict {
+            Some(indices) => indices
                 .iter()
-            .map(|&idx| tokens[idx].as_ref())
+                .filter_map(|&idx| tokens.get(idx).map(|cow| cow.as_ref()))
-            .collect()
+                .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) 包含校验
+            // 史诗级加速点：如果是 None，说明没限制，根本不进入分支，直接放行！
-            // 注意：去掉原本错误的 `!`
+            // 只有当有具体限制（Some）时，才去跑 4-5 次 CPU 寄存器级别的二分查找
-            if has_any_match {
+            if let Some(ref indices) = self.charset_restrict {
-                if !valid_indices.contains(&u_idx) {
+                if indices.binary_search(&u_idx).is_err() {
-                    continue; // 不在有效字符集内，安全跳过
+                    continue;
                 }
             }