refactor: 重构 OcrBuilder 过滤机制，实现零克隆、高性能的局部白名单

- 将有状态的有效索引缓存 `valid_indices` 从全局 `Charset` 剥离至局部 `OcrBuilder`
- 解码函数 `ctc_decode_to_string` 内部复用策略计算，通过 `has_any_match` 布尔开关实现全量免检短路加速
- 优化内存分配，根据 `CharsetRestrict` 策略动态精准计算 `HashSet` 初始容量，规避大词表空置浪费
- 增强鲁棒性，在策略与字库完全无交集时，自动触发智能降级，一键恢复全量识别并保持接口自洽

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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
-    pub charset: Vec<String>,
+    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<String>) -> Result<Self, anyhow::Error> {
+    pub fn new(model_path: String, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error> {
         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<Vec<ColorRange>>,
     /// 字符集范围
-    charset_range: Option<CharsetRange>,
+    charset_restrict: Option<CharsetRestrict>,
 }
 
 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<String, anyhow::Error> {
+    pub fn run(&self) -> anyhow::Result<String> {
         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::<i64>()?;
-        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<Vec<i64>> {
+    fn extract_indices_from_tensor(&self, raw_tensor: &Tensor) -> anyhow::Result<Vec<i64>> {
         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<usize> {
+        let (_, valid_indices) = self.valid_indices();
+        valid_indices
+    }
+
+    fn valid_indices(&self) -> (bool, HashSet<usize>) {
+        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
     }
 }