CNWei
913ff4d884
- 新增 Other变体以及构造函数new - 剥离图像预处理中的 Base64 相关错误至业务层处理 - 引入强类型 `LogitsDimensionMismatch` 替代不便匹配的字符串错误 - 优化 `normalize_ocr_logits` 的转换流程，兼顾零拷贝性能与精细化报错 - 优化 全库错误处理
74 lines
2.9 KiB
Rust
74 lines
2.9 KiB
Rust
use crate::loader::ModelLoader;
|
|
use anyhow::Context;
|
|
use ddddocr_core::error::{DdddError, Result, TensorErrorReason};
|
|
use ddddocr_core::{DetEngine, DetOutput, InferenceEngine};
|
|
use ndarray::Ix3;
|
|
use std::path::Path;
|
|
use tract_onnx::prelude::{Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tvec};
|
|
#[derive(Debug)]
|
|
pub struct DetSession {
|
|
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
|
}
|
|
|
|
impl DetSession {
|
|
pub fn new<P>(model_path: P) -> Result<Self>
|
|
where
|
|
P: AsRef<Path>,
|
|
{
|
|
let session = ModelLoader::model_for_path(&model_path)?.session;
|
|
Ok(Self { session })
|
|
}
|
|
|
|
pub fn model_from_bytes(model_bytes: &[u8]) -> Result<Self> {
|
|
let session = ModelLoader::model_from_bytes(model_bytes)?.session;
|
|
Ok(Self { session })
|
|
}
|
|
// pub fn inference(&self, tensor: Tensor) -> anyhow::Result<Tensor> {
|
|
// // tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
|
|
// // let result = self.ocr.run(tvec!(tensor.into()))?;
|
|
// let mut result = self
|
|
// .session
|
|
// .run(tvec!(tensor.into()))
|
|
// .context("执行模型推理失败")?;
|
|
// println!("模型输出原始数据: {:?}", result);
|
|
// Ok(result.swap_remove(0).into_tensor())
|
|
// }
|
|
}
|
|
|
|
impl InferenceEngine for DetSession {
|
|
type Output = DetOutput; // 明确绑定 OCR 小枚举
|
|
fn inference(&self, input_array: ndarray::Array4<f32>) -> Result<Self::Output> {
|
|
// 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(|_| {
|
|
DdddError::Inference(TensorErrorReason::EngineError(
|
|
"执行模型推理失败".to_string(),
|
|
))
|
|
})?;
|
|
println!("模型输出原始数据: {:?}", result);
|
|
// Ok(result.swap_remove(0).into_tensor())
|
|
let raw_tensor = result.swap_remove(0).into_tensor();
|
|
let array_d = raw_tensor.into_array::<f32>().map_err(|_| {
|
|
DdddError::Inference(TensorErrorReason::EngineError(
|
|
"Tract 实体张量无法转换为 ndarray::ArrayD".to_string(),
|
|
))
|
|
})?;
|
|
// 提前利用克隆(Clone)备份好当前未转维度前的真实 shape (Vec<usize>)
|
|
let actual_shape = array_d.shape().to_vec();
|
|
|
|
let array3 = array_d.into_dimensionality::<Ix3>().map_err(|_| {
|
|
DdddError::Inference(TensorErrorReason::TensorDimensionMismatch {
|
|
expected: "3D 检测矩阵 [Batch, Box_Count, Box_Attributes]".to_string(),
|
|
actual: actual_shape, // 优雅降维失败时动态捕获
|
|
})
|
|
})?;
|
|
Ok(DetOutput::Detection(array3))
|
|
|
|
// 在引擎内部消化掉 DatumType 强耦合
|
|
}
|
|
}
|
|
|
|
impl DetEngine for DetSession {}
|