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

refactor(ocr): 拆分字符集限制枚举并引入声明式多路组合子宏

Browse Source 
- 将原本臃肿的 CharsetRestrict 拆分为 CharRestrict（内容过滤）和 IdRestrict（索引过滤），实现职责解耦
 - 引入 any_of! 声明式宏实现无 Box、无堆内存分配的栈上多路组合，规避孤儿规则
 - 完善 estimated_capacity 容量预估函数，实现真正的 O(1) 精准内存开辟
This commit is contained in:
CNWei
2026-06-16 09:37:15 +08:00
parent b7146831f7
commit 189f2bd697
4 changed files with 111 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
src/charset.rs
View File

@@ -517,83 +517,109 @@ pub const CHARSET_BETA: &[&str] = &[
pub const CHARSET_OLD: &[&str] = &["", "笤", "谴", "膀", "荔"];
pub fn get_default_charset() -> Vec<String> {
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<String>),
/// 完美对应 Python 传入 int 时的行为：截取并只保留模型字符集前 N 个字符
TopN(usize),
/// 核心组合子：满足左边或右边任意一个条件即可（即 A + B 的并集逻辑）
/// 使用 Box 打破 Rust 编译期对递归枚举的无限大小限制
Or(Box<CharsetRestrict>, Box<CharsetRestrict>),
}
impl From<i32> 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<usize>),
IdList(Vec<usize>),
}
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<String>，内部实现真正的零拷贝
pub tokens: Vec<Cow<'static, str>>,
pub tokens: Vec<Cow<'static, str>>,
// 反向查找表，保证字符转索引为 O(1)
pub char_to_idx: HashMap<Cow<'static, str>, 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(),)
}
}

4
src/lib.rs
View File

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

2
src/model_metadata.rs
View File

@@ -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;

49
src/models/ocr.rs
View File

@@ -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<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;