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 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, // 既包含 Fixed 静态维度,也包含 Dynamic 动态符号 // pub data_type: TensorElementType, // 对应 Python 的 type // } // // /// 最终返回的模型完整信息 // #[derive(Debug, Clone)] // pub struct ModelInfo { // pub inputs: Vec, // pub outputs: Vec, // /// 硬件执行提供者(采用 Option 兼容不同底层的推理引擎) // pub providers: Option>, // } pub struct OcrRuntime { pub session: Session, pub metadata: ModelMetadata, } impl OcrRuntime { pub fn new(session: Session, metadata: ModelMetadata) -> Self { Self { session, metadata } } } impl OcrEngine for OcrRuntime { fn metadata(&self) -> &ModelMetadata { &self.metadata } } impl InferenceEngine for OcrRuntime { type Output = OcrOutput; /// 对应 Python 的 _inference fn inference(&self, input_array: ndarray::Array4) -> Result { // tract 的 run 会返回一个 Vec,我们通常只需要第一个输出 // 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()))?; 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]; match raw_value.dtype().tensor_type().unwrap() { TensorElementType::Int64 => { let (array_d, slice) = raw_value .try_extract_tensor::() .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::>(); 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::() .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::() .map_err(|_| TensorError::Engine("Tract 无法获取 f32 内存视图".to_string()))?; // 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗 let shape_vec: Vec = 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 Info for OcrRuntime { fn input_info(&self) -> Result> { todo!() } fn output_info(&self) -> Result> { todo!() } fn model_info(&self) -> Result { todo!() } }