CNWei
ea7fb43a14
- 移除 核心层与 tract/Tensor 的强耦合，前/后处理全线转用标准 ndarray - 针对 OCR 与目标检测（Det）分别设计独立的强类型输出小枚举（OcrOutput/DetOutput） - 利用 Trait 关联类型（Associated Type）InferenceEngine,OcrEngine,DetEngine 统一接口，实现多后端解耦 - 引入 thiserror 库，建立完备的强类型错误处理机制（DdddError/Result） - 完成项目结构初拆，剥离为 ddddocr-core 和 ddddocr-tract
201 lines
6.6 KiB
Rust
201 lines
6.6 KiB
Rust
use anyhow::{anyhow, Result};
|
||
use serde::Deserialize;
|
||
use std::borrow::Cow;
|
||
use std::collections::HashMap;
|
||
|
||
// ==========================================
|
||
// 3. 字符集核心结构体 (重命名为 Charset)
|
||
// ==========================================
|
||
#[derive(Debug, Clone)]
|
||
pub struct Charset {
|
||
// 使用 Cow 统一静态切片和动态读取的 Vec<String>，内部实现真正的零拷贝
|
||
pub tokens: Vec<Cow<'static, str>>,
|
||
// 反向查找表，保证字符转索引为 O(1)
|
||
pub char_to_idx: HashMap<Cow<'static, str>, usize>,
|
||
// 当前处于激活状态的有效索引缓存 (用于 CTC 解码前的过滤加速)
|
||
// pub 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);
|
||
}
|
||
|
||
Self {
|
||
tokens,
|
||
char_to_idx,
|
||
}
|
||
}
|
||
|
||
// --- 业务策略方法 ---
|
||
|
||
/// 将字符转为索引，不存在返回 -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 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()
|
||
}
|
||
}
|
||
|
||
// ==========================================
|
||
// 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: {}", self.size(),)
|
||
}
|
||
}
|
||
|
||
|
||
|
||
// =====================================================================
|
||
// 1. 辅助定义的枚举与结构体
|
||
// =====================================================================
|
||
|
||
#[derive(Debug, Clone, Copy, Deserialize)]
|
||
#[serde(rename_all = "snake_case")] // 支持 json 中写 "zero_to_one" 或 "minus_one_to_one"
|
||
pub enum Normalization {
|
||
/// 映射到 [0.0, 1.0] -> pixel / 255.0
|
||
ZeroToOne,
|
||
/// 映射到 [-1.0, 1.0] -> (pixel / 255.0 - 0.5) / 0.5
|
||
MinusOneToOne,
|
||
}
|
||
|
||
impl Normalization {
|
||
/// 统一归一化计算逻辑
|
||
#[inline(always)]
|
||
pub fn normalize(&self, pixel: f32) -> f32 {
|
||
match self {
|
||
Normalization::ZeroToOne => pixel / 255.0,
|
||
Normalization::MinusOneToOne => (pixel / 255.0 - 0.5) / 0.5,
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 图像缩放策略枚举
|
||
#[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,
|
||
/// 新增：允许在配置文件中指定归一化策略。
|
||
/// 使用 serde(default) 可以在不配置时提供一个默认值（比如默认 ZeroToOne）
|
||
#[serde(default = "default_normalization")]
|
||
normalization: Normalization,
|
||
}
|
||
fn default_normalization() -> Normalization {
|
||
Normalization::ZeroToOne
|
||
}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct ModelMetadata {
|
||
/// 字符集管理器
|
||
pub charset: Charset,
|
||
/// 是否为单字识别模型
|
||
pub word: bool,
|
||
/// 预处理的缩放策略
|
||
pub resize: Resize,
|
||
/// 图像通道数 (1 或 3)
|
||
pub channel: u8,
|
||
/// 新增：传递给核心业务使用的归一化配置
|
||
pub normalization: Normalization,
|
||
}
|
||
|
||
impl ModelMetadata {
|
||
// --- 优雅的工厂模式构造器 ---
|
||
/// 通用的静态切片转换构造器
|
||
pub fn from_static_slice(
|
||
slice: &[&'static str],
|
||
word: bool,
|
||
resize: Resize,
|
||
channel: u8,
|
||
normalization: Normalization,
|
||
) -> Self {
|
||
let tokens: Vec<Cow<'static, str>> = slice.iter().map(|&s| Cow::Borrowed(s)).collect();
|
||
Self {
|
||
charset: Charset::new(tokens),
|
||
word,
|
||
resize,
|
||
channel,
|
||
normalization,
|
||
}
|
||
}
|
||
pub fn from_json_str(json_str: &str) -> Result<Self> {
|
||
let dto: ModelMetadataDto = serde_json::from_str(json_str)
|
||
.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,
|
||
normalization: dto.normalization,
|
||
})
|
||
}
|
||
/// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载）
|
||
pub fn from_json_bytes(bytes: &[u8]) -> Result<Self> {
|
||
let json_str = std::str::from_utf8(bytes)
|
||
.map_err(|e| anyhow!("JSON 字节流不是合法的 UTF-8 编码: {}", e))?;
|
||
Self::from_json_str(json_str)
|
||
}
|
||
}
|