From 189f2bd697f412fea90ec56a96022d8b5fc99bed Mon Sep 17 00:00:00 2001 From: CNWei Date: Tue, 16 Jun 2026 09:37:15 +0800 Subject: [PATCH] =?UTF-8?q?refactor(ocr):=20=E6=8B=86=E5=88=86=E5=AD=97?= =?UTF-8?q?=E7=AC=A6=E9=9B=86=E9=99=90=E5=88=B6=E6=9E=9A=E4=B8=BE=E5=B9=B6?= =?UTF-8?q?=E5=BC=95=E5=85=A5=E5=A3=B0=E6=98=8E=E5=BC=8F=E5=A4=9A=E8=B7=AF?= =?UTF-8?q?=E7=BB=84=E5=90=88=E5=AD=90=E5=AE=8F=20=20-=20=E5=B0=86?= =?UTF-8?q?=E5=8E=9F=E6=9C=AC=E8=87=83=E8=82=BF=E7=9A=84=20CharsetRestrict?= =?UTF-8?q?=20=E6=8B=86=E5=88=86=E4=B8=BA=20CharRestrict=EF=BC=88=E5=86=85?= =?UTF-8?q?=E5=AE=B9=E8=BF=87=E6=BB=A4=EF=BC=89=E5=92=8C=20IdRestrict?= =?UTF-8?q?=EF=BC=88=E7=B4=A2=E5=BC=95=E8=BF=87=E6=BB=A4=EF=BC=89=EF=BC=8C?= =?UTF-8?q?=E5=AE=9E=E7=8E=B0=E8=81=8C=E8=B4=A3=E8=A7=A3=E8=80=A6=20=20-?= =?UTF-8?q?=20=E5=BC=95=E5=85=A5=20any=5Fof!=20=E5=A3=B0=E6=98=8E=E5=BC=8F?= =?UTF-8?q?=E5=AE=8F=E5=AE=9E=E7=8E=B0=E6=97=A0=20Box=E3=80=81=E6=97=A0?= =?UTF-8?q?=E5=A0=86=E5=86=85=E5=AD=98=E5=88=86=E9=85=8D=E7=9A=84=E6=A0=88?= =?UTF-8?q?=E4=B8=8A=E5=A4=9A=E8=B7=AF=E7=BB=84=E5=90=88=EF=BC=8C=E8=A7=84?= =?UTF-8?q?=E9=81=BF=E5=AD=A4=E5=84=BF=E8=A7=84=E5=88=99=20=20-=20?= =?UTF-8?q?=E5=AE=8C=E5=96=84=20estimated=5Fcapacity=20=E5=AE=B9=E9=87=8F?= =?UTF-8?q?=E9=A2=84=E4=BC=B0=E5=87=BD=E6=95=B0=EF=BC=8C=E5=AE=9E=E7=8E=B0?= =?UTF-8?q?=E7=9C=9F=E6=AD=A3=E7=9A=84=20O(1)=20=E7=B2=BE=E5=87=86?= =?UTF-8?q?=E5=86=85=E5=AD=98=E5=BC=80=E8=BE=9F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/charset.rs | 158 ++++++++++++++++++++++++------------------ src/lib.rs | 4 +- src/model_metadata.rs | 2 +- src/models/ocr.rs | 49 +++++-------- 4 files changed, 111 insertions(+), 102 deletions(-) diff --git a/src/charset.rs b/src/charset.rs index 52a5805..6a88a71 100644 --- a/src/charset.rs +++ b/src/charset.rs @@ -517,83 +517,109 @@ pub const CHARSET_BETA: &[&str] = &[ pub const CHARSET_OLD: &[&str] = &["", "笤", "谴", "膀", "荔"]; -pub fn get_default_charset() -> Vec { - CHARSET_BETA.iter().map(|&s| s.to_string()).collect() -} - use std::borrow::Cow; use std::collections::{HashMap, HashSet}; -use std::ops::{Add, Deref}; -// 字符集范围类型 -/// 字符集范围限制组合子枚举 +/// 字符集范围限制枚举 +pub struct ValidationCtx<'a> { + pub text: &'a str, // 当前 Token 的文本内容 + pub token_id: usize, // 当前 Token 的 ID 索引 +} + +/// 统一的约束接口 +pub trait TokenFilter { + fn matches(&self, ctx: &ValidationCtx) -> bool; + /// 预估容量提示，帮助精准开辟 Vec 内存 + fn estimated_capacity(&self) -> usize { + 128 + } +} + #[derive(Debug, Clone, PartialEq, Eq)] -pub enum CharsetRestrict { - /// 纯整数 0-9 +pub enum CharRestrict { Digit, - - /// 纯小写字母 a-z Lowercase, - - /// 纯大写字母 A-Z Uppercase, - - // /// 过滤模式：删除所有 ASCII 字母和数字（通常用于仅保留汉字、特殊标点） - // // ExcludeAlphanumeric, - // /// 自定义单字字符集，例如 "0123456789+-x/=" - // Single(String), - /// 直接设置完整的 Token 白名单（支持多字 Token），例如 vec!["html".to_string()] CustomList(Vec), - /// 完美对应 Python 传入 int 时的行为：截取并只保留模型字符集前 N 个字符 - TopN(usize), - /// 核心组合子：满足左边或右边任意一个条件即可（即 A + B 的并集逻辑） - /// 使用 Box 打破 Rust 编译期对递归枚举的无限大小限制 - Or(Box, Box), } -impl From for CharsetRestrict { - fn from(value: i32) -> Self { - match value { - 0 => Self::Digit, - 1 => Self::Lowercase, - 2 => Self::Uppercase, - // 3 => Self::LowercaseUppercase, - // 4 => Self::LowercaseDigit, - // 5 => Self::UppercaseDigit, - // 6 => Self::LowercaseUppercaseDigit, - // 7 => Self::DefaultCharsetLowercaseUppercaseDigit, - _ => panic!("invalid charset range: {}", value), - } - } -} -impl CharsetRestrict { - /// 💡 辅助构造函数：直接在源头把用户的长字符串切碎，伪装成基础积木 - pub fn from_chars(custom_str: &str) -> Self { - let tokens = custom_str.chars().map(|c| c.to_string()).collect(); - CharsetRestrict::CustomList(tokens) - } - // 内部递归收集器：利用硬编码切片快速无损展开 - pub(crate) fn matches(&self, s: &str) -> bool { + +impl TokenFilter for CharRestrict { + fn matches(&self, ctx: &ValidationCtx) -> bool { match self { - CharsetRestrict::Digit => s.len() == 1 && s.as_bytes()[0].is_ascii_digit(), - 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), + Self::Digit => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_digit(), + Self::Lowercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_lowercase(), + Self::Uppercase => ctx.text.len() == 1 && ctx.text.as_bytes()[0].is_ascii_uppercase(), + Self::CustomList(vec) => vec.iter().any(|t| t == ctx.text), + } + } + fn estimated_capacity(&self) -> usize { + match self { + Self::Digit => 16, + Self::Lowercase | Self::Uppercase => 32, + Self::CustomList(vec) => vec.len() + 1, } } } -// ===================================================================== -// 5. 优雅的魔法：重载 + 运算符 (实现 std::ops::Add) -// ===================================================================== -/// 支持 `CharsetRestrict::Digit + CharsetRestrict::Lowercase` -impl Add for CharsetRestrict { - type Output = Self; +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum IdRestrict { + TopN(usize), + IdRange(std::ops::Range), + IdList(Vec), +} - fn add(self, rhs: Self) -> Self::Output { - CharsetRestrict::Or(Box::new(self), Box::new(rhs)) +impl TokenFilter for IdRestrict { + fn matches(&self, ctx: &ValidationCtx) -> bool { + match self { + Self::TopN(n) => ctx.token_id < *n, + Self::IdRange(range) => range.contains(&ctx.token_id), + Self::IdList(vec) => vec.contains(&ctx.token_id), + } } + fn estimated_capacity(&self) -> usize { + match self { + Self::TopN(n) => *n + 1, + // 2. IdRange：标准标准库 Range 的长度 + // 注意：因为范围可能是 1000..2000，它的 len() 返回的是 usize + Self::IdRange(range) => range.len() + 1, + // 3. IdList：Vec 里的元素个数 + Self::IdList(vec) => vec.len() + 1, + } + } +} + +/// 多路“或”逻辑组合子（支持 N 个规则无缝并集） +pub struct MultiOrRestrict<'a> { + pub filters: Vec<&'a dyn TokenFilter>, +} + +impl<'a> TokenFilter for MultiOrRestrict<'a> { + fn matches(&self, ctx: &ValidationCtx) -> bool { + // 核心高阶函数：只要有一个过滤器命中，该 Token 即可放行 + self.filters.iter().any(|f| f.matches(ctx)) + } + + fn estimated_capacity(&self) -> usize { + // 将所有过滤器的预估容量累加，作为最终容量参考 + self.filters.iter().map(|f| f.estimated_capacity()).sum() + } +} +// ===================================================================== +// 声明式宏：替代 `+` 运算符，解决组合扩展痛苦 +// ===================================================================== +#[macro_export] +macro_rules! any_of { + // 场景 A：如果用户只传了一个规则，免去构建 Vec 的开销，直接返回其引用 + ($only:expr) => { + &$only as &dyn $crate::TokenFilter + }; + + // 场景 B：如果用户传入了多个规则，自动织成一张静态组合网 + ($($filter:expr),+ $(,)?) => { + &$crate::MultiOrRestrict { + filters: vec![ $( &$filter as &dyn $crate::TokenFilter ),+ ] + } + }; } // ========================================== @@ -602,7 +628,7 @@ impl Add for CharsetRestrict { #[derive(Debug, Clone)] pub struct Charset { // 使用 Cow 统一静态切片和动态读取的 Vec，内部实现真正的零拷贝 - pub tokens: Vec>, + pub tokens: Vec>, // 反向查找表，保证字符转索引为 O(1) pub char_to_idx: HashMap, usize>, // 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速) @@ -641,14 +667,12 @@ impl Charset { 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() } pub fn size(&self) -> usize { self.tokens.len() } - } // ========================================== @@ -656,10 +680,6 @@ impl Charset { // ========================================== impl std::fmt::Display for Charset { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "Charset [Total Size: {}", - self.size(), - ) + write!(f, "Charset [Total Size: {}", self.size(),) } } diff --git a/src/lib.rs b/src/lib.rs index 29e3d8c..2781c36 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, CharsetRestrict}; +use crate::charset::{ CharRestrict}; use crate::models::ocr::ColorRange; use models::det::Det; use models::loader::ModelSession; @@ -97,7 +97,7 @@ impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { // Runtime::Ocr(s) => s.predict(img).run(), - Runtime::Ocr(s) => s.builder().charset_restrict(&CharsetRestrict::Digit).predict(img), + Runtime::Ocr(s) => s.builder().charset_restrict(&CharRestrict::Digit).predict(img), Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), } } diff --git a/src/model_metadata.rs b/src/model_metadata.rs index 24bd590..dec971a 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -1,4 +1,4 @@ -use crate::charset::{CHARSET_BETA, CHARSET_OLD, Charset, CharsetRestrict}; +use crate::charset::{CHARSET_BETA, CHARSET_OLD, Charset}; use anyhow::{Result, anyhow}; use serde::Deserialize; use std::borrow::Cow; diff --git a/src/models/ocr.rs b/src/models/ocr.rs index 6071fe5..f9e10ba 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -1,4 +1,4 @@ -use crate::charset::CharsetRestrict; +use crate::charset::{TokenFilter, ValidationCtx}; use crate::model_metadata::ModelMetadata; use crate::models::base::ModelArgs; use crate::models::loader::{ModelLoader, ModelSession, ModelType}; @@ -158,45 +158,34 @@ impl<'a> OcrBuilder<'a> { self.color_filter_custom_ranges = Some(value); self } - pub fn charset_restrict(mut self, restrict: &CharsetRestrict) -> Self { + pub fn charset_restrict(mut self, restrict: &dyn TokenFilter) -> Self { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; // let mut temp_indices = Vec::new(); let mut has_any_match = false; - 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 黄金保底值 - }; + let estimated_capacity = restrict.estimated_capacity(); // 精准开辟内存，完美避开 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); + for (idx, token) in tokens.iter().enumerate() { + let token_str = token.as_ref(); + // CTC Blank 空字符串无条件放行，其余交给超高性能的 matches + if token_str.is_empty() || idx == 0 { + temp_indices.push(idx); + }; + // 组装无拷贝上下文 + let ctx = ValidationCtx { + text: token_str, + token_id: idx, + }; + + if restrict.matches(&ctx) { + temp_indices.push(idx); 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() { - temp_indices.push(idx); - } else if restrict.matches(token_str) { - temp_indices.push(idx); - has_any_match = true; - } - } } + // self.charset_restrict = Some(restrict); // 终极防御：如果除了 Blank 外什么都没匹配上，退化恢复为 None（全量识别） if !has_any_match { @@ -266,7 +255,7 @@ impl<'a> OcrBuilder<'a> { // 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;