refactor(core): 提炼公共类型
- 将 AxisDim、TensorInfo 等公共类型下沉至 ddddocr_core::types - 项目结构优化
This commit is contained in:
@@ -1,70 +1,69 @@
|
||||
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 ort::inputs;
|
||||
use ort::value::{TensorElementType, TensorRef};
|
||||
use std::sync::Mutex;
|
||||
// 引入核心层的统一错误类型
|
||||
/// 明确命名为 AxisDim,代表模型某一个轴的维度特征
|
||||
#[derive(Clone, PartialEq, Eq)]
|
||||
pub enum AxisDim {
|
||||
/// 静态固定维度(如通道数固定为 1,高度固定为 64)
|
||||
Static(usize),
|
||||
/// 动态符号维度(如宽度是动态的 "image_width")
|
||||
Dynamic(String),
|
||||
}
|
||||
// #[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),
|
||||
}
|
||||
}
|
||||
}
|
||||
// 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: TensorElementType, // 对应 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 {
|
||||
// #[derive(Debug, Clone)]
|
||||
// pub struct TensorInfo {
|
||||
// pub name: String,
|
||||
// pub shape: Vec<AxisDim>, // 既包含 Fixed 静态维度,也包含 Dynamic 动态符号
|
||||
// pub data_type: TensorElementType, // 对应 Python 的 type
|
||||
// }
|
||||
//
|
||||
// /// 最终返回的模型完整信息
|
||||
// #[derive(Debug, Clone)]
|
||||
// pub struct ModelInfo {
|
||||
// pub inputs: Vec<TensorInfo>,
|
||||
// pub outputs: Vec<TensorInfo>,
|
||||
// /// 硬件执行提供者(采用 Option 兼容不同底层的推理引擎)
|
||||
// pub providers: Option<Vec<String>>,
|
||||
// }
|
||||
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 }
|
||||
}
|
||||
}
|
||||
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> {
|
||||
@@ -76,18 +75,6 @@ impl InferenceEngine for OcrSession {
|
||||
.lock()
|
||||
.map_err(|_| TensorError::Engine("获取 Session 锁失败 (Poisoned)".to_string()))?;
|
||||
|
||||
// // 2. 获取输入节点名称
|
||||
// let input_name = session_guard
|
||||
// .inputs()
|
||||
// .first()
|
||||
// .map(|i| i.name())
|
||||
// .unwrap_or("input");
|
||||
//
|
||||
// // 3. 在 session_guard (&mut Session) 上调用 run
|
||||
// let outputs = session_guard
|
||||
// .run(inputs![TensorRef::from_array_view(&input_array).map_err(|e| TensorError::Engine(format!("构建输入失败: {e}")))? )
|
||||
// .map_err(|e| TensorError::Engine(format!("执行模型推理失败: {e}")))?;
|
||||
|
||||
let result = session_guard
|
||||
.run(inputs![TensorRef::from_array_view(&input_array).map_err(
|
||||
|e| TensorError::Engine(format!("构建输入失败: {e}"))
|
||||
@@ -97,10 +84,7 @@ impl InferenceEngine for OcrSession {
|
||||
println!("模型输出原始数据: {:?}", result);
|
||||
// Ok(result.swap_remove(0).into_tensor())
|
||||
let raw_value = &result[0];
|
||||
// let dtype = raw_tensor
|
||||
// .dtype();
|
||||
// .map_err(|e| TensorError::Engine(format!("无法读取输出数据类型: {e}")))?;
|
||||
// 在引擎内部消化掉 DatumType 强耦合
|
||||
|
||||
match raw_value.dtype().tensor_type().unwrap() {
|
||||
TensorElementType::Int64 => {
|
||||
let (array_d, slice) = raw_value
|
||||
@@ -108,7 +92,7 @@ impl InferenceEngine for OcrSession {
|
||||
.map_err(|_| TensorError::Engine("Tract 无法获取 i64 内存视图".to_string()))?;
|
||||
// .context("Tract 无法获取 i64 内存视图")?;
|
||||
|
||||
// 🌟 提前提取真实维度
|
||||
// 提前提取真实维度
|
||||
let actual_shape = array_d
|
||||
.to_vec()
|
||||
.iter()
|
||||
@@ -151,75 +135,16 @@ impl InferenceEngine for OcrSession {
|
||||
}
|
||||
}
|
||||
}
|
||||
// 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
|
||||
// }
|
||||
// }
|
||||
impl Info for OcrRuntime {
|
||||
fn input_info(&self) -> Result<Vec<TensorInfo>> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn output_info(&self) -> Result<Vec<TensorInfo>> {
|
||||
todo!()
|
||||
}
|
||||
|
||||
fn model_info(&self) -> Result<ModelInfo> {
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user