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

refactor(core): 提炼公共类型

Browse Source 
- 将 AxisDim、TensorInfo 等公共类型下沉至 ddddocr_core::types
- 项目结构优化
This commit is contained in:
CNWei
2026-07-30 16:55:58 +08:00
parent 7d159c5702
commit a3c4614574
22 changed files with 349 additions and 408 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
ddddocr-tract/src/det/session.rs
View File

@@ -1,22 +1,22 @@
use crate::types::Session;
use ddddocr_core::error::{Result, TensorError};
use ddddocr_core::{DetEngine, DetOutput, InferenceEngine};
use ddddocr_core::{ DetOutput};
use ddddocr_core::traits::{DetEngine, InferenceEngine};
use ndarray::Ix3;
// use tract_onnx::prelude::{tvec, IntoTensor, Tensor};
use tract_onnx::prelude::*;
#[derive(Debug)]
pub struct DetSession {
pub struct DetRuntime {
pub session: Session,
}
impl DetSession {
impl DetRuntime {
pub fn new(session: Session) -> Self {
Self { session }
}
}
impl InferenceEngine for DetSession {
impl InferenceEngine for DetRuntime {
type Output = DetOutput; // 明确绑定 OCR 小枚举
fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output, TensorError> {
// tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
@@ -49,4 +49,4 @@ impl InferenceEngine for DetSession {
}
}
impl DetEngine for DetSession {}
impl DetEngine for DetRuntime {}

10
ddddocr-tract/src/lib.rs
View File

@@ -1,9 +1,11 @@
mod det;
mod error;
pub mod loader;
mod ocr;
mod types;
mod error;
pub use ddddocr_core::{SlideResult, Slider,OcrBuilder};
pub use det::session::DetSession;
pub use ocr::session::OcrSession;
pub use ddddocr_core::{
DetectionResult, Detector, ModelMetadata, Normalization, Ocr, OcrBuilder, SlideResult, Slider,
};
pub use det::session::DetRuntime;
pub use ocr::session::OcrRuntime;

4
ddddocr-tract/src/loader.rs
View File

@@ -3,5 +3,5 @@ mod metadata;
mod model;
pub use error::{Error, ParseError, Result};
pub use metadata::{ModelMetadataDto, NormalizationDto, TractModelMetadata};
pub use model::TractModelLoader;
pub use metadata::{ModelMetadataDto, NormalizationDto, Metadata};
pub use model::ModelLoader;

4
ddddocr-tract/src/loader/metadata.rs
View File

@@ -42,7 +42,7 @@ fn default_normalization() -> NormalizationDto {
}
/// Tract 专属扩展trait 或 工具函数
pub trait TractModelMetadata: Sized {
pub trait Metadata: Sized {
fn from_json_str(json_str: &str) -> Result<Self>;
/// 机制 2：从内存字节流加载（极大地方便 include_bytes! 或网络下载）
fn from_json_bytes(bytes: &[u8]) -> Result<Self> {
@@ -50,7 +50,7 @@ pub trait TractModelMetadata: Sized {
Self::from_json_str(json_str)
}
}
impl TractModelMetadata for ModelMetadata {
impl Metadata for ModelMetadata {
// --- 优雅的工厂模式构造器 ---
fn from_json_str(json_str: &str) -> Result<ModelMetadata> {
let dto: ModelMetadataDto = serde_json::from_str(json_str)?;

20
ddddocr-tract/src/loader/model.rs
View File

@@ -1,26 +1,18 @@
use crate::loader::error::{Error, ParseError, Result};
use crate::types::Session;
use ddddocr_core::ModelBuilder;
use ddddocr_core::traits::Loader;
use std::io::Cursor;
use tract_linalg::multithread::{Executor, set_default_executor};
use tract_onnx::onnx;
use tract_onnx::prelude::*;
pub struct TractModelLoader;
impl TractModelLoader {
/// 获取针对 Tract 后端的链式构建器
pub fn builder() -> TractModelBuilder {
TractModelBuilder::default()
}
}
/// Tract 专用的链式构建器
#[derive(Debug, Clone, Default)]
pub struct TractModelBuilder {
pub struct ModelLoader {
num_threads: Option<usize>,
}
impl TractModelBuilder {
impl ModelLoader {
/// 可选扩展：设置 CPU 线程数（不提供任何 GPU 相关的 API）
pub fn num_threads(mut self, threads: usize) -> Self {
self.num_threads = Some(threads);
@@ -41,10 +33,10 @@ impl TractModelBuilder {
}
}
impl ModelBuilder for TractModelBuilder {
impl Loader for ModelLoader {
type Session = Session;
type Error = Error;
fn model_for_path<P>(&self, model_path: P) -> Result<Session>
fn build_for_path<P>(&self, model_path: P) -> Result<Session>
where
P: AsRef<std::path::Path>,
{
@@ -64,7 +56,7 @@ impl ModelBuilder for TractModelBuilder {
Ok(session)
}
/// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用）
fn model_from_bytes(&self, model_bytes: &[u8]) -> Result<Session> {
fn build_from_bytes(&self, model_bytes: &[u8]) -> Result<Session> {
self.setup_tract_threads();
// 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read）
let mut cursor = Cursor::new(model_bytes);

170
ddddocr-tract/src/ocr/session.rs
View File

@@ -1,69 +1,80 @@
use crate::types::Session;
use ddddocr_core::ModelMetadata;
use ddddocr_core::OcrOutput;
use ddddocr_core::error::{DdddError, Result, TensorError};
use ddddocr_core::traits::{InferenceEngine, Info, OcrEngine};
use ddddocr_core::types::{AxisDim, ModelInfo, TensorInfo, TensorType};
use ddddocr_core::utils::normalize_ocr_logits;
use ddddocr_core::{InferenceEngine, OcrEngine, OcrOutput};
use tract_onnx::prelude::{DatumType, OutletId, ShapeFact, TypedModel};
use tract_onnx::prelude::{IntoTensor, Tensor, tvec};
// 引入核心层的统一错误类型
/// 明确命名为 AxisDim，代表模型某一个轴的维度特征
#[derive(Clone, PartialEq, Eq)]
pub enum AxisDim {
/// 静态固定维度（如通道数固定为 1，高度固定为 64）
Static(usize),
/// 动态符号维度（如宽度是动态的 "image_width"）
Dynamic(String),
}
impl AxisDim {
/// 便捷方法：判断是否为动态维度
pub fn is_dynamic(&self) -> bool {
matches!(self, AxisDim::Dynamic(_))
}
}
/// 自定义 Debug 格式化输出，彻底融化套娃外壳，保证日志干净漂亮
impl std::fmt::Debug for AxisDim {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
AxisDim::Static(size) => write!(f, "{}", size),
AxisDim::Dynamic(expr) => write!(f, "Dynamic(\"{}\")", expr),
}
}
}
/// 模拟 Python 的 input_info 和 output_info 结构
#[derive(Debug, Clone)]
pub struct TensorInfo {
pub name: String,
pub shape: Vec<AxisDim>, // 既包含 Fixed 静态维度，也包含 Dynamic 动态符号
pub data_type: DatumType, // 对应 Python 的 type
}
/// 最终返回的模型完整信息
#[derive(Debug, Clone)]
pub struct ModelInfo {
pub inputs: Vec<TensorInfo>,
pub outputs: Vec<TensorInfo>,
/// 硬件执行提供者（采用 Option 兼容不同底层的推理引擎）
pub providers: Option<Vec<String>>,
}
pub struct OcrSession {
pub struct OcrRuntime {
pub session: Session,
pub model_metadata: ModelMetadata,
pub metadata: ModelMetadata,
}
impl OcrSession {
pub fn new(session: Session, model_metadata: ModelMetadata) -> Self {
Self {
session,
model_metadata,
}
impl OcrRuntime {
pub fn new(session: Session, metadata: ModelMetadata) -> Self {
Self { session, metadata }
}
/// 获取模型输入的节点信息列表
/// 提取出来的公共转换逻辑：将一组 OutletId 解析为 TensorInfo 列表
fn resolve_tensors(&self, model: &TypedModel, outlets: &[OutletId]) -> Result<Vec<TensorInfo>> {
outlets
.iter()
.map(|&outlet_id| {
let fact = model.outlet_fact(outlet_id).map_err(DdddError::new)?;
// .map_err(|e| {
// DdddError::InternalError(format!("解析节点 Fact 失败: {:?}", e))
// })?;
let shape = self.resolve_shape(&fact.shape);
let node_name = model.node(outlet_id.node).name.clone();
let tensor_type = match fact.datum_type {
DatumType::F32 => TensorType::F32,
DatumType::I64 => TensorType::I64,
_ => TensorType::Other,
};
Ok(TensorInfo {
name: node_name,
shape,
tensor_type,
})
})
.collect() // 函数式声明：自动传播第一处发生的错误
}
/// 安全还原 Tract 维度至 Vec<AxisDim>
fn resolve_shape(&self, shape_fact: &ShapeFact) -> Vec<AxisDim> {
let tract_shape = shape_fact.to_tvec();
let resolved = tract_shape
.iter()
.map(|dim| {
// 防御性编程：必须同时满足能够转换为 i64 且 大于等于 0
if let Ok(size) = dim.to_i64() {
if size >= 0 {
AxisDim::Static(size as usize)
} else {
// 如果 ONNX 导出时某些动态维度被标记为了 -1，安全地作为动态符号捕获
AxisDim::Dynamic(dim.to_string())
}
} else {
AxisDim::Dynamic(dim.to_string())
}
})
.collect();
resolved
}
}
impl OcrEngine for OcrSession {
impl OcrEngine for OcrRuntime {
fn metadata(&self) -> &ModelMetadata {
&self.model_metadata
&self.metadata
}
}
impl InferenceEngine for OcrSession {
impl InferenceEngine for OcrRuntime {
type Output = OcrOutput;
/// 对应 Python 的 _inference
fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output, TensorError> {
@@ -115,16 +126,15 @@ impl InferenceEngine for OcrSession {
}
}
}
impl OcrSession {
/// 获取模型输入的节点信息列表
pub fn input_info(&self) -> Result<Vec<TensorInfo>> {
impl Info for OcrRuntime {
fn input_info(&self) -> Result<Vec<TensorInfo>> {
let model = self.session.model();
let outlets = model.input_outlets().map_err(DdddError::new)?;
self.resolve_tensors(model, outlets)
}
/// 获取模型输出的节点信息列表
pub fn output_info(&self) -> Result<Vec<TensorInfo>> {
fn output_info(&self) -> Result<Vec<TensorInfo>> {
let model = self.session.model();
let outlets = model.output_outlets().map_err(DdddError::new)?;
self.resolve_tensors(model, outlets)
@@ -133,57 +143,11 @@ impl OcrSession {
/// 获取模型详细元数据信息（对标 Python ddddocr 的 get_model_info）
/// 完美包容 [1, 1, 64, image_width] 这样的变长图像模型
/// 获取模型详细元数据信息（代码更紧凑、优雅）
pub fn model_info(&self) -> Result<ModelInfo> {
fn model_info(&self) -> Result<ModelInfo> {
Ok(ModelInfo {
inputs: self.input_info()?,
outputs: self.output_info()?,
providers: None,
})
}
/// 提取出来的公共转换逻辑：将一组 OutletId 解析为 TensorInfo 列表
fn resolve_tensors(&self, model: &TypedModel, outlets: &[OutletId]) -> Result<Vec<TensorInfo>> {
outlets
.iter()
.map(|&outlet_id| {
let fact = model.outlet_fact(outlet_id).map_err(DdddError::new)?;
// .map_err(|e| {
// DdddError::InternalError(format!("解析节点 Fact 失败: {:?}", e))
// })?;
let shape = self.resolve_shape(&fact.shape);
let node_name = model.node(outlet_id.node).name.clone();
Ok(TensorInfo {
name: node_name,
shape,
data_type: fact.datum_type,
})
})
.collect() // 函数式声明：自动传播第一处发生的错误
}
/// 安全还原 Tract 维度至 Vec<AxisDim>
fn resolve_shape(&self, shape_fact: &ShapeFact) -> Vec<AxisDim> {
let tract_shape = shape_fact.to_tvec();
let resolved = tract_shape
.iter()
.map(|dim| {
// 防御性编程：必须同时满足能够转换为 i64 且 大于等于 0
if let Ok(size) = dim.to_i64() {
if size >= 0 {
AxisDim::Static(size as usize)
} else {
// 如果 ONNX 导出时某些动态维度被标记为了 -1，安全地作为动态符号捕获
AxisDim::Dynamic(dim.to_string())
}
} else {
AxisDim::Dynamic(dim.to_string())
}
})
.collect();
resolved
}
}

43
ddddocr-tract/tests/ocr_test.rs
View File

@@ -1,8 +1,9 @@
use anyhow::Context;
use ddddocr_core::{DetectionResult, ModelBuilder};
use ddddocr_core::{Detector, ModelMetadata, Normalization, Slider};
use ddddocr_core::traits::Loader;
use ddddocr_tract::{DetectionResult, Ocr};
use ddddocr_tract::{Detector, ModelMetadata, Normalization, Slider};
// 假设你的包名是这个
use ddddocr_tract::{DetSession, OcrBuilder, OcrSession};
use ddddocr_tract::{DetRuntime, OcrRuntime};
use image::{DynamicImage, ImageBuffer, Luma, Rgb};
use std::fs;
use std::path::Path;
@@ -11,7 +12,7 @@ mod char_slice;
use char_slice::CHARSET_BETA;
use ddddocr_core::Resize;
use ddddocr_tract::loader::{ TractModelLoader};
use ddddocr_tract::loader::ModelLoader as TractModelLoader;
fn load_image<P: AsRef<Path>>(path: P) -> anyhow::Result<image::DynamicImage> {
// 1. 先将泛型转为具体的 &Path 引用
@@ -104,9 +105,10 @@ fn save_rust_result(result: &ImageBuffer<Luma<f32>, Vec<f32>>, filename: &str) {
}
#[test]
fn test_full_classification() {
let model = TractModelLoader::builder()
.model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx")
let session = TractModelLoader::default()
.build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx")
.expect("模型加载失败");
let metadata = ModelMetadata::from_static_slice(
CHARSET_BETA,
false,
@@ -115,18 +117,18 @@ fn test_full_classification() {
Normalization::MinusOneToOne,
);
// 1. 初始化模型
let ocr = OcrSession::new(model, metadata);
let ocr_runtime = OcrRuntime::new(session, metadata);
// 2. 加载测试图片
let img =
image::open("D:/CNWei/CNW/Rust/ddddocr-rs/samples/code2.png").expect("测试图片不存在");
// 3. 执行识别
// let result = Ocr::new(&ocr)
// let result = Ocr::new(&ocr_runtime)
// .predict(&img)
// .expect("识别过程出错")
// .into_text();
let result = OcrBuilder::new()
.build(&ocr)
let result = Ocr::builder()
.runner(&ocr_runtime)
.predict(&img)
.expect("识别过程出错")
.into_text();
@@ -136,11 +138,10 @@ fn test_full_classification() {
}
#[test]
fn test_det_load() -> anyhow::Result<()> {
let det_model =
TractModelLoader::builder()
.model_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")
.expect("模型加载失败");
let det = DetSession::new(det_model);
let det_model = TractModelLoader::default()
.build_for_path("D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_det.onnx")
.expect("模型加载失败");
let det = DetRuntime::new(det_model);
let image_path = "D:/CNWei/CNW/Rust/ddddocr-rs/samples/det1.png";
let image_bytes =
fs::read(image_path).map_err(|e| anyhow::anyhow!("无法读取图片 {}: {}", image_path, e))?;
@@ -240,11 +241,11 @@ fn test_real_slide_comparison() {
#[test]
fn test_resolve_shape_logic_direct() {
// 创建一个哑 ModelLoader 实例（session 用不上，因为我们直接测私有方法）
let loader = TractModelLoader::builder()
.model_for_path(
// "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx",
"D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx",
)
.expect("建立测试模型图失败");
let loader = TractModelLoader::default()
.build_for_path(
// "D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_sml2h3_f32.onnx",
"D:\\CNWei\\CNW\\Rust\\ddddocr-rs\\models\\common_huashi666_i64.onnx",
)
.expect("建立测试模型图失败");
println!("{:?}", loader.model().inputs);
}