feat(model): 新增 ModelLoader 链式构建 API 及 ORT GPU/Tract 多线程配置
- 在 ddddocr-core 中定义 ModelBuilder Trait 及其错误类型 - ddddocr-ort 支持 use_gpu、device_id 及 num_threads 链式配置与 CUDA 硬件加速 - ddddocr-tract 基于 multithread-mm 特性支持 CPU 线程数控制 - 支持基于 tract-linalg 配置推理线程数，显式引入 tract-linalg 的 multithread-mm 特性，控制 GEMM 算子并发 - 优化线程池加载策略，适配 Tokio 异步及 CLI 等多场景
This commit is contained in:
225
ddddocr-ort/src/ocr/session.rs
Normal file
225
ddddocr-ort/src/ocr/session.rs
Normal file
@@ -0,0 +1,225 @@
|
||||
use crate::types::Session;
|
||||
use ddddocr_core::ModelMetadata;
|
||||
use ddddocr_core::error::{DdddError, Result, TensorError};
|
||||
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),
|
||||
}
|
||||
|
||||
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 {
|
||||
pub session: Session,
|
||||
pub model_metadata: ModelMetadata,
|
||||
}
|
||||
impl OcrSession {
|
||||
pub fn new(session: Session, model_metadata: ModelMetadata) -> Self {
|
||||
Self {
|
||||
session,
|
||||
model_metadata,
|
||||
}
|
||||
}
|
||||
}
|
||||
impl OcrEngine for OcrSession {
|
||||
fn metadata(&self) -> &ModelMetadata {
|
||||
&self.model_metadata
|
||||
}
|
||||
}
|
||||
impl InferenceEngine for OcrSession {
|
||||
type Output = OcrOutput;
|
||||
/// 对应 Python 的 _inference
|
||||
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);
|
||||
let mut session_guard = self
|
||||
.session
|
||||
.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}"))
|
||||
)?])
|
||||
.map_err(|e| TensorError::Engine(format!("执行模型推理失败: {e}")))?;
|
||||
// .context("执行模型推理失败")?;
|
||||
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
|
||||
.try_extract_tensor::<i64>()
|
||||
.map_err(|_| TensorError::Engine("Tract 无法获取 i64 内存视图".to_string()))?;
|
||||
// .context("Tract 无法获取 i64 内存视图")?;
|
||||
|
||||
// 🌟 提前提取真实维度
|
||||
let actual_shape = array_d
|
||||
.to_vec()
|
||||
.iter()
|
||||
.map(|v| *v as usize)
|
||||
.collect::<Vec<usize>>();
|
||||
|
||||
let view = ndarray::ArrayViewD::from_shape(actual_shape.as_slice(), slice)
|
||||
.map_err(|_| TensorError::Engine("构建 ndarray ArrayViewD 失败".to_string()))?;
|
||||
// 转成标准的 Array1 传给 core
|
||||
let array1 = view
|
||||
.to_owned()
|
||||
.into_dimensionality::<ndarray::Ix1>()
|
||||
.map_err(|_| TensorError::DimensionMismatch {
|
||||
expected: "1D 字符索引静态矩阵".to_string(),
|
||||
actual: actual_shape,
|
||||
})?;
|
||||
Ok(OcrOutput::Indices(array1))
|
||||
}
|
||||
TensorElementType::Float32 => {
|
||||
let shape = raw_value.shape();
|
||||
println!("模型输出shape数据: {:?}", shape);
|
||||
// raw_tensor.to_plain_array_view()
|
||||
let (shape_ref, slice) = raw_value
|
||||
.try_extract_tensor::<f32>()
|
||||
.map_err(|_| TensorError::Engine("Tract 无法获取 f32 内存视图".to_string()))?;
|
||||
// 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗
|
||||
let shape_vec: Vec<usize> =
|
||||
shape_ref.to_vec().iter().map(|v| *v as usize).collect();
|
||||
let shape_vec_slice = shape_vec.as_slice();
|
||||
|
||||
let view = ndarray::ArrayViewD::from_shape(shape_vec_slice, slice)
|
||||
.map_err(|_| TensorError::Engine("构建 ndarray ArrayViewD 失败".to_string()))?;
|
||||
|
||||
normalize_ocr_logits(view, shape_vec_slice)
|
||||
}
|
||||
_ => Err(
|
||||
// anyhow::anyhow!("不支持的模型输出数据类型: {:?}",raw_tensor.datum_type())
|
||||
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