54 lines
1.8 KiB
Rust
54 lines
1.8 KiB
Rust
use crate::models::ocr::model_metadata::ModelMetadata;
|
||
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
|
||
use anyhow::Context;
|
||
use anyhow::Result;
|
||
use std::path::Path;
|
||
use tract_onnx::prelude::{tvec, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp};
|
||
|
||
pub struct OcrSession {
|
||
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
||
pub model_metadata: ModelMetadata,
|
||
}
|
||
impl ModelSession for OcrSession {
|
||
fn get_model_type(&self) -> ModelType {
|
||
todo!("使用thiserror作为错误处理的库,thiserror 专门用于开发库（Library）");
|
||
}
|
||
fn desc(&self) -> String {
|
||
"Ocr Model 加载成功".to_string()
|
||
}
|
||
}
|
||
impl OcrSession {
|
||
pub fn new<P>(model_path: P, model_metadata: ModelMetadata) -> Result<Self, anyhow::Error>
|
||
where
|
||
P: AsRef<Path>,
|
||
{
|
||
let session = ModelLoader::model_for_path(model_path)?.session;
|
||
Ok(Self {
|
||
session,
|
||
model_metadata,
|
||
})
|
||
}
|
||
|
||
pub fn model_from_bytes(
|
||
model_bytes: &[u8],
|
||
model_metadata: ModelMetadata,
|
||
) -> Result<Self, anyhow::Error> {
|
||
let session = ModelLoader::model_from_bytes(model_bytes)?.session;
|
||
Ok(Self {
|
||
session,
|
||
model_metadata,
|
||
})
|
||
}
|
||
/// 对应 Python 的 _inference
|
||
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())
|
||
}
|
||
}
|