refactor(load,tract):将 ModelMetadata JSON 加载逻辑解耦至 ddddocr-tract, 优化 Error 枚举结构与错误透传
- 在 load 模块中精简 Error 与 Result 别名定义 - 增加 ParseError 子类型区分路径与字节流加载失败 - 支持通过 #[from] 自动转换 Tract 引擎底层错误 - 移出 core 中的 serde 依赖，保持核心库纯洁 - 在 tract 中实现 TractModelMetadata 扩展 trait 加载解析配置
This commit is contained in:
@@ -1,35 +1,61 @@
|
||||
use crate::loader::ModelLoader;
|
||||
use anyhow::Context;
|
||||
use crate::loader::ModelMetadata;
|
||||
use crate::types::Session;
|
||||
use ddddocr_core::error::{DdddError, Result, TensorError};
|
||||
use ddddocr_core::utils::normalize_ocr_logits;
|
||||
use ddddocr_core::{InferenceEngine, ModelMetadata, OcrEngine, OcrOutput};
|
||||
use ndarray::s;
|
||||
use std::path::Path;
|
||||
use tract_onnx::prelude::DatumType;
|
||||
use tract_onnx::prelude::{Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tvec};
|
||||
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 session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
||||
pub session: Session,
|
||||
pub model_metadata: ModelMetadata,
|
||||
}
|
||||
impl OcrSession {
|
||||
pub fn new<P>(model_path: P, model_metadata: ModelMetadata) -> Result<Self>
|
||||
where
|
||||
P: AsRef<Path>,
|
||||
{
|
||||
let session = ModelLoader::model_for_path(model_path)?.session;
|
||||
Ok(Self {
|
||||
pub fn new(session: Session, model_metadata: ModelMetadata) -> Self {
|
||||
Self {
|
||||
session,
|
||||
model_metadata,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn model_from_bytes(model_bytes: &[u8], model_metadata: ModelMetadata) -> Result<Self> {
|
||||
let session = ModelLoader::model_from_bytes(model_bytes)?.session;
|
||||
Ok(Self {
|
||||
session,
|
||||
model_metadata,
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
impl OcrEngine for OcrSession {
|
||||
@@ -40,7 +66,7 @@ impl OcrEngine for OcrSession {
|
||||
impl InferenceEngine for OcrSession {
|
||||
type Output = OcrOutput;
|
||||
/// 对应 Python 的 _inference
|
||||
fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output> {
|
||||
fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output, TensorError> {
|
||||
// tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
|
||||
// let result = self.ocr.run(tvec!(tensor.into()))?;
|
||||
let tensor = Tensor::from(input_array);
|
||||
@@ -48,12 +74,8 @@ impl InferenceEngine for OcrSession {
|
||||
let mut result = self
|
||||
.session
|
||||
.run(tvec!(tensor.into()))
|
||||
.map_err(|_| {
|
||||
DdddError::Inference(TensorError::EngineError(
|
||||
"执行模型推理失败".to_string(),
|
||||
))
|
||||
})?;
|
||||
// .context("执行模型推理失败")?;
|
||||
.map_err(|_| TensorError::Engine("执行模型推理失败".to_string()))?;
|
||||
// .context("执行模型推理失败")?;
|
||||
println!("模型输出原始数据: {:?}", result);
|
||||
// Ok(result.swap_remove(0).into_tensor())
|
||||
let raw_tensor = result.swap_remove(0).into_tensor();
|
||||
@@ -62,23 +84,17 @@ impl InferenceEngine for OcrSession {
|
||||
DatumType::I64 => {
|
||||
let array_d = raw_tensor
|
||||
.into_array::<i64>()
|
||||
.map_err(|_| {
|
||||
DdddError::Inference(TensorErrorReason::EngineError(
|
||||
"Tract 无法获取 i64 内存视图".to_string(),
|
||||
))
|
||||
})?;
|
||||
// .context("Tract 无法获取 i64 内存视图")?;
|
||||
.map_err(|_| TensorError::Engine("Tract 无法获取 i64 内存视图".to_string()))?;
|
||||
// .context("Tract 无法获取 i64 内存视图")?;
|
||||
// 🌟 提前提取真实维度
|
||||
let actual_shape = array_d.shape().to_vec();
|
||||
// 转成标准的 Array1 传给 core
|
||||
let array1 = array_d
|
||||
.to_owned()
|
||||
.into_dimensionality::<ndarray::Ix1>()
|
||||
.map_err(|_| {
|
||||
DdddError::Inference(TensorErrorReason::TensorDimensionMismatch {
|
||||
expected: "1D 字符索引静态矩阵".to_string(),
|
||||
actual: actual_shape,
|
||||
})
|
||||
.map_err(|_| TensorError::DimensionMismatch {
|
||||
expected: "1D 字符索引静态矩阵".to_string(),
|
||||
actual: actual_shape,
|
||||
})?;
|
||||
Ok(OcrOutput::Indices(array1))
|
||||
}
|
||||
@@ -87,18 +103,86 @@ impl InferenceEngine for OcrSession {
|
||||
println!("模型输出shape数据: {:?}", shape);
|
||||
let view = raw_tensor
|
||||
.to_array_view::<f32>()
|
||||
.map_err(|_| {
|
||||
DdddError::Inference(TensorErrorReason::EngineError(
|
||||
"Tract 无法获取 f32 内存视图".to_string(),
|
||||
))
|
||||
})?;
|
||||
.map_err(|_| TensorError::Engine("Tract 无法获取 f32 内存视图".to_string()))?;
|
||||
// 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗
|
||||
normalize_ocr_logits(view, shape)
|
||||
}
|
||||
_ => Err(
|
||||
// anyhow::anyhow!("不支持的模型输出数据类型: {:?}",raw_tensor.datum_type())
|
||||
DdddError::Inference(TensorErrorReason::UnknownOutputFormat)
|
||||
TensorError::UnknownOutputFormat,
|
||||
),
|
||||
}
|
||||
}
|
||||
}
|
||||
impl OcrSession {
|
||||
/// 获取模型输入的节点信息列表
|
||||
pub 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>> {
|
||||
let model = self.session.model();
|
||||
let outlets = model.output_outlets().map_err(DdddError::new)?;
|
||||
self.resolve_tensors(model, outlets)
|
||||
}
|
||||
|
||||
/// 获取模型详细元数据信息（对标 Python ddddocr 的 get_model_info）
|
||||
/// 完美包容 [1, 1, 64, image_width] 这样的变长图像模型
|
||||
/// 获取模型详细元数据信息（代码更紧凑、优雅）
|
||||
pub 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
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user