ddddocr-rs/ddddocr-ort/src/ocr/session.rs

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<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 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<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()))?;

        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::<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 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!()
    }
}