CNWei/ddddocr-rs
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 将有状态的有效索引缓存 `valid_indices` 从全局 `Charset` 剥离至局部 `OcrBuilder`
- 解码函数 `ctc_decode_to_string` 内部复用策略计算，通过 `has_any_match` 布尔开关实现全量免检短路加速
- 优化内存分配，根据 `CharsetRestrict` 策略动态精准计算 `HashSet` 初始容量，规避大词表空置浪费
- 增强鲁棒性，在策略与字库完全无交集时，自动触发智能降级，一键恢复全量识别并保持接口自洽
This commit is contained in:
CNWei
2026-06-09 18:12:02 +08:00
parent 15ce068025
commit 0c96fbedbf
5 changed files with 133 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files

60
src/charset.rs
View File

@@ -600,11 +600,11 @@ impl Add for CharsetRestrict {
#[derive(Debug, Clone)]
pub struct Charset {
// 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
tokens: Vec<Cow<'static, str>>,
pub tokens: Vec<Cow<'static, str>>,
// 反向查找表，保证字符转索引为 O(1)
char_to_idx: HashMap<Cow<'static, str>, usize>,
pub char_to_idx: HashMap<Cow<'static, str>, usize>,
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
valid_indices: HashSet<usize>,
// pub valid_indices: HashSet<usize>,
}
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) {
@@ -673,20 +639,6 @@ impl Charset {
self.tokens.get(index).map(|cow| cow.as_ref())
}
/// 获取有效字符索引列表 (用于外部验证或过滤)
pub fn get_valid_indices(&self) -> Vec<usize> {
let mut indices: Vec<usize> = 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()
@@ -695,9 +647,6 @@ impl Charset {
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()
)
}
}

11
src/lib.rs
View File

@@ -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<String>,
model_metadata: ModelMetadata,
},
}
impl ModelSpec {
@@ -61,9 +62,9 @@ impl DdddOcrBuilder {
}
/// 设置自定义 OCR 路径
pub fn custom_ocr(mut self, path: String, charset: Vec<String>) -> 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 })

148
src/models/ocr.rs
View File

@@ -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) {
// 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; // 不在有效字符集内，安全跳过
}
}
// 5. 字符映射
if let Some(char_str) = tokens.get(u_idx) {
res.push_str(char_str);
} else {
// 保护逻辑：如果模型预测的索引超出了字符集范围
eprintln!("警告: 预测索引 {} 超出字符集范围", u_idx);
}
}
}
prev_idx = idx;
}
println!("最终识别出的验证码是: {}", res);
res
}
}

2
tests/ocr_test.rs
View File

@@ -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("识别过程出错");