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

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 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: 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 result = self
            .session
            .run(tvec!(tensor.into()))
            .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();
        // 在引擎内部消化掉 DatumType 强耦合
        match raw_tensor.datum_type() {
            DatumType::I64 => {
                let array_d = raw_tensor
                    .into_plain_array::<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(|_| TensorError::DimensionMismatch {
                        expected: "1D 字符索引静态矩阵".to_string(),
                        actual: actual_shape,
                    })?;
                Ok(OcrOutput::Indices(array1))
            }
            DatumType::F32 => {
                let shape = raw_tensor.shape();
                println!("模型输出shape数据: {:?}", shape);
                // raw_tensor.to_plain_array_view()
                let view = raw_tensor
                    .to_plain_array_view::<f32>()
                    .map_err(|_| TensorError::Engine("Tract 无法获取 f32 内存视图".to_string()))?;
                // 1. 极其纯粹的、无拷贝的多维 Shape 压扁清洗
                normalize_ocr_logits(view, shape)
            }
            _ => 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
    }
}