feat: 字符集限制枚举优化与核心解码器对接
- 新增 model_metadata.rs - 优化 charset.rs - 其他优化
This commit is contained in:
198
src/charset.rs
198
src/charset.rs
@@ -514,6 +514,202 @@ 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};
|
||||
// 字符集范围类型
|
||||
/// 字符集范围限制组合子枚举
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
pub enum CharsetRestrict {
|
||||
/// 纯整数 0-9
|
||||
Digit,
|
||||
|
||||
/// 纯小写字母 a-z
|
||||
Lowercase,
|
||||
|
||||
/// 纯大写字母 A-Z
|
||||
Uppercase,
|
||||
|
||||
// /// 过滤模式：删除所有 ASCII 字母和数字（通常用于仅保留汉字、特殊标点）
|
||||
// // ExcludeAlphanumeric,
|
||||
// /// 自定义单字字符集，例如 "0123456789+-x/="
|
||||
// Single(String),
|
||||
/// 直接设置完整的 Token 白名单（支持多字 Token），例如 vec!["html".to_string()]
|
||||
CustomList(Vec<String>),
|
||||
|
||||
/// 核心组合子：满足左边或右边任意一个条件即可（即 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 {
|
||||
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::Or(left, right) => left.matches(s) || right.matches(s),
|
||||
}
|
||||
}
|
||||
}
|
||||
// =====================================================================
|
||||
// 5. 优雅的魔法：重载 + 运算符 (实现 std::ops::Add)
|
||||
// =====================================================================
|
||||
|
||||
/// 支持 `CharsetRestrict::Digit + CharsetRestrict::Lowercase`
|
||||
impl Add for CharsetRestrict {
|
||||
type Output = Self;
|
||||
|
||||
fn add(self, rhs: Self) -> Self::Output {
|
||||
CharsetRestrict::Or(Box::new(self), Box::new(rhs))
|
||||
}
|
||||
}
|
||||
|
||||
// ==========================================
|
||||
// 3. 字符集核心结构体 (重命名为 Charset)
|
||||
// ==========================================
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct Charset {
|
||||
// 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
|
||||
tokens: Vec<Cow<'static, str>>,
|
||||
// 反向查找表，保证字符转索引为 O(1)
|
||||
char_to_idx: HashMap<Cow<'static, str>, usize>,
|
||||
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
|
||||
valid_indices: HashSet<usize>,
|
||||
}
|
||||
|
||||
impl Charset {
|
||||
// 内部底层统一收拢构造
|
||||
pub fn new(tokens: Vec<Cow<'static, str>>) -> Self {
|
||||
let mut char_to_idx = HashMap::with_capacity(tokens.len());
|
||||
for (idx, token) in tokens.iter().enumerate() {
|
||||
char_to_idx.entry(token.clone()).or_insert(idx);
|
||||
// 如果字符集有重复，保留第一个遇到的索引 (符合 Python .index 逻辑)
|
||||
// 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) {
|
||||
idx as i32
|
||||
} else {
|
||||
-1
|
||||
}
|
||||
}
|
||||
|
||||
/// 将索引转为字符引用，零拷贝。若越界返回 None
|
||||
pub fn index_to_char_ref(&self, index: usize) -> Option<&str> {
|
||||
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()
|
||||
}
|
||||
pub fn size(&self) -> usize {
|
||||
self.tokens.len()
|
||||
}
|
||||
|
||||
pub fn valid_size(&self) -> usize {
|
||||
self.valid_indices.len()
|
||||
}
|
||||
}
|
||||
|
||||
// ==========================================
|
||||
// 4. 标准 Display 接口实现 (对应 __str__)
|
||||
// ==========================================
|
||||
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: {}]",
|
||||
self.size(),
|
||||
self.valid_size()
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,6 +2,7 @@ mod charset;
|
||||
|
||||
pub mod models;
|
||||
pub mod utils;
|
||||
mod model_metadata;
|
||||
|
||||
use anyhow::Result;
|
||||
use image::DynamicImage;
|
||||
@@ -94,7 +95,7 @@ impl Display for DdddOcr {
|
||||
impl DdddOcr {
|
||||
pub fn classification(&self, img: &DynamicImage) -> Result<String> {
|
||||
match &self.runtime {
|
||||
Runtime::Ocr(s) => s.predict(img, false),
|
||||
Runtime::Ocr(s) => s.predict(img).run(),
|
||||
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")),
|
||||
}
|
||||
}
|
||||
|
||||
122
src/model_metadata.rs
Normal file
122
src/model_metadata.rs
Normal file
@@ -0,0 +1,122 @@
|
||||
use crate::charset::{CHARSET_BETA, CHARSET_OLD, Charset, CharsetRestrict};
|
||||
use anyhow::{Result, anyhow};
|
||||
use serde::Deserialize;
|
||||
use std::borrow::Cow;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::fs::File;
|
||||
use std::io::Read;
|
||||
use std::path::Path;
|
||||
// =====================================================================
|
||||
// 1. 辅助定义的枚举与结构体
|
||||
// =====================================================================
|
||||
|
||||
/// 图像缩放策略枚举
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Resize {
|
||||
/// 固定宽高，例如 (64, 64)
|
||||
Fixed(u32, u32),
|
||||
/// 高度固定，宽度根据原始比例动态计算（对应 Python 的 [-1, H]）
|
||||
DynamicWidth(u32),
|
||||
/// 单字识别的正方形切图（对应 Python 的 word 为 True 且 [-1, H]）
|
||||
Square(u32),
|
||||
}
|
||||
|
||||
/// 仅用于反序列化 JSON 的中间临时结构体（DTO）
|
||||
#[derive(Deserialize)]
|
||||
struct ModelMetadataDto {
|
||||
charset: Vec<String>,
|
||||
word: bool,
|
||||
#[serde(alias = "image")]
|
||||
resize: Vec<i32>,
|
||||
channel: u8,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ModelMetadata {
|
||||
/// 字符集管理器
|
||||
pub charset: Charset,
|
||||
/// 是否为单字识别模型
|
||||
pub word: bool,
|
||||
/// 预处理的缩放策略
|
||||
pub resize: Resize,
|
||||
/// 图像通道数 (1 或 3)
|
||||
pub channel: u8,
|
||||
}
|
||||
|
||||
impl ModelMetadata {
|
||||
// --- 优雅的工厂模式构造器 ---
|
||||
/// 从预设的旧版字符集创建
|
||||
pub fn from_builtin_old() -> Self {
|
||||
Self::from_static_slice(CHARSET_OLD, false, Resize::Fixed(64, 64), 1)
|
||||
}
|
||||
|
||||
/// 从预设的 Beta 版字符集创建
|
||||
pub fn from_builtin_beta() -> Self {
|
||||
Self::from_static_slice(CHARSET_BETA, false, Resize::Fixed(64, 64), 1)
|
||||
}
|
||||
|
||||
/// 通用的静态切片转换构造器
|
||||
pub fn from_static_slice(
|
||||
slice: &[&'static str],
|
||||
word: bool,
|
||||
resize: Resize,
|
||||
channel: u8,
|
||||
) -> Self {
|
||||
let tokens: Vec<Cow<'static, str>> = slice.iter().map(|&s| Cow::Borrowed(s)).collect();
|
||||
Self {
|
||||
charset: Charset::new(tokens),
|
||||
word,
|
||||
resize,
|
||||
channel,
|
||||
}
|
||||
}
|
||||
|
||||
/// 从外部外部 JSON 文件动态加载字符集
|
||||
pub fn from_json_file<P: AsRef<Path>>(path: P) -> Result<Self> {
|
||||
let path = path.as_ref();
|
||||
if !path.exists() {
|
||||
return Err(anyhow!("模型元数据配置文件不存在: {:?}", path));
|
||||
}
|
||||
|
||||
let mut file = File::open(path)?;
|
||||
let mut content = String::new();
|
||||
file.read_to_string(&mut content)?;
|
||||
|
||||
let dto: ModelMetadataDto = serde_json::from_str(&content)
|
||||
.map_err(|e| anyhow!("JSON 反序列化失败，请检查字段是否完整: {}", e))?;
|
||||
|
||||
// 1. 将 DTO 的字符串数组转化为强类型的 Charset
|
||||
let tokens: Vec<Cow<'static, str>> =
|
||||
dto.charset.into_iter().map(|s| Cow::Owned(s)).collect();
|
||||
let charset = Charset::new(tokens);
|
||||
|
||||
// 2. 解析 resize 策略（重现 Python 的复杂条件判断）
|
||||
if dto.resize.len() != 2 {
|
||||
return Err(anyhow!(
|
||||
"'resize (or image)' 字段必须是包含两个元素的数组，例如 [-1, 64]"
|
||||
));
|
||||
}
|
||||
let r0 = dto.resize[0];
|
||||
let r1 = dto.resize[1];
|
||||
|
||||
let resize = if r0 == -1 {
|
||||
if dto.word {
|
||||
// 如果 word 为 true，且包含 -1，Python 里是 resize 为 (r1, r1) 的正方形
|
||||
Resize::Square(r1 as u32)
|
||||
} else {
|
||||
// 如果 word 为 false，且包含 -1，Python 里是高度固定为 r1，宽度按原图比例缩放
|
||||
Resize::DynamicWidth(r1 as u32)
|
||||
}
|
||||
} else {
|
||||
// 正常的固定宽高
|
||||
Resize::Fixed(r0 as u32, r1 as u32)
|
||||
};
|
||||
|
||||
Ok(Self {
|
||||
charset,
|
||||
word: dto.word,
|
||||
resize,
|
||||
channel: dto.channel,
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -8,92 +8,13 @@ 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));
|
||||
|
||||
// 字符集范围类型
|
||||
#[derive(Debug, Clone)]
|
||||
pub enum CharsetRange {
|
||||
All, // 所有字符
|
||||
Digit, // 数字
|
||||
Letter, // 字母
|
||||
Alphanumeric, // 字母数字
|
||||
Single(String), // 单字符串
|
||||
Multiple(Vec<String>), // 多个字符串
|
||||
Range(char, char), // 字符范围
|
||||
Custom(Vec<char>), // 自定义字符列表
|
||||
}
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct PredictArgs {
|
||||
/// 是否修复PNG格式问题
|
||||
pub png_fix: bool,
|
||||
/// 是否返回概率信息
|
||||
pub probability: bool,
|
||||
/// 颜色过滤：保留的颜色列表
|
||||
pub color_filter_colors: Option<Vec<String>>,
|
||||
/// 颜色过滤：自定义RGB范围
|
||||
pub color_filter_custom_ranges: Option<Vec<ColorRange>>,
|
||||
/// 字符集范围
|
||||
pub charset_range: Option<CharsetRange>,
|
||||
}
|
||||
|
||||
impl Default for PredictArgs {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
png_fix: false,
|
||||
probability: false,
|
||||
color_filter_colors: None,
|
||||
color_filter_custom_ranges: None,
|
||||
charset_range: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl PredictArgs {
|
||||
pub fn new() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
// Builder 模式方法
|
||||
pub fn png_fix(mut self, enabled: bool) -> Self {
|
||||
self.png_fix = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn probability(mut self, enabled: bool) -> Self {
|
||||
self.probability = enabled;
|
||||
self
|
||||
}
|
||||
|
||||
pub fn color_filter_colors(mut self, colors: Vec<String>) -> Self {
|
||||
self.color_filter_colors = Some(colors);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn color_filter_custom_ranges(mut self, ranges: Vec<ColorRange>) -> Self {
|
||||
self.color_filter_custom_ranges = Some(ranges);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn charset_range(mut self, range: CharsetRange) -> Self {
|
||||
self.charset_range = Some(range);
|
||||
self
|
||||
}
|
||||
|
||||
// 便捷构造方法
|
||||
pub fn quick() -> Self {
|
||||
Self::default()
|
||||
}
|
||||
|
||||
pub fn with_probability() -> Self {
|
||||
Self::default().probability(true)
|
||||
}
|
||||
|
||||
pub fn with_png_fix() -> Self {
|
||||
Self::default().png_fix(true)
|
||||
}
|
||||
}
|
||||
pub struct Ocr {
|
||||
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
||||
pub charset: Vec<String>,
|
||||
@@ -111,28 +32,38 @@ impl Ocr {
|
||||
let session = ModelLoader::load_model(&model_path)?.session;
|
||||
Ok(Self { session, charset })
|
||||
}
|
||||
pub fn task<'a>(&'a self, image: &'a DynamicImage) -> OcrTask {
|
||||
OcrTask::new(self, image)
|
||||
pub fn predict<'a>(&'a self, image: &'a DynamicImage) -> OcrBuilder<'a> {
|
||||
OcrBuilder::new(self, image)
|
||||
}
|
||||
}
|
||||
|
||||
pub struct OcrTask<'a> {
|
||||
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_range: Option<CharsetRange>,
|
||||
}
|
||||
|
||||
impl<'a> OcrTask<'a> {
|
||||
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_range: None
|
||||
}
|
||||
}
|
||||
pub fn png_fix(mut self, value: bool) -> Self {
|
||||
@@ -147,9 +78,13 @@ impl<'a> OcrTask<'a> {
|
||||
self.color_filter_custom_ranges = Some(value);
|
||||
self
|
||||
}
|
||||
pub fn charset_range(mut self, range: CharsetRange) -> Self {
|
||||
self.charset_range = Some(range);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn predict(&self, image: &DynamicImage, png_fix: bool) -> Result<String, anyhow::Error> {
|
||||
let tensor = self.preprocess_image(image, png_fix)?;
|
||||
pub fn run(&self) -> Result<String, anyhow::Error> {
|
||||
let tensor = self.preprocess_image(self.image, self.png_fix)?;
|
||||
//
|
||||
// let result = self.session.run(tvec!(tensor.into()))?;
|
||||
// // 3. 解析结果
|
||||
|
||||
Reference in New Issue
Block a user