49 lines
1.3 KiB
Rust
49 lines
1.3 KiB
Rust
use anyhow::Context;
|
||
use std::io::Cursor;
|
||
use tract_onnx::onnx;
|
||
use tract_onnx::prelude::*;
|
||
|
||
/// OCR 模型：包含路径和字符集
|
||
|
||
pub enum ModelType {
|
||
Ocr,
|
||
Det,
|
||
Custom,
|
||
}
|
||
// 定义统一的 trait
|
||
pub trait ModelSession {
|
||
fn get_model_type(&self) -> ModelType;
|
||
fn desc(&self) -> String;
|
||
}
|
||
|
||
pub struct ModelLoader {
|
||
pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
|
||
}
|
||
|
||
impl ModelLoader {
|
||
pub fn model_for_path<P>(model_path: P) -> anyhow::Result<Self>
|
||
where
|
||
P: AsRef<std::path::Path>,
|
||
{
|
||
let session = onnx()
|
||
.model_for_path(model_path)
|
||
.with_context(|| "加载 ONNX 模型失败，请检查路径是否正确")?
|
||
.into_optimized()?
|
||
.into_runnable()?;
|
||
Ok(Self { session })
|
||
}
|
||
/// 策略 B：从内存字节流加载模型（配合 include_bytes! 使用）
|
||
pub fn model_from_bytes(model_bytes: &[u8]) -> anyhow::Result<Self> {
|
||
// 使用 std::io::Cursor 将 &[u8] 包装为可读的流（实现 std::io::Read）
|
||
let mut cursor = Cursor::new(model_bytes);
|
||
|
||
let session = onnx()
|
||
.model_for_read(&mut cursor)
|
||
.with_context(|| "从内存字节流解析 ONNX 模型失败")?
|
||
.into_optimized()?
|
||
.into_runnable()?;
|
||
|
||
Ok(Self { session })
|
||
}
|
||
}
|