refactor(slide,det): 优化项目结构，移除不必要的逻辑
- 优化 项目结构，移除不必要的逻辑
This commit is contained in:
245
src/lib.rs
245
src/lib.rs
@@ -1,138 +1,141 @@
|
||||
mod charset;
|
||||
|
||||
mod error;
|
||||
mod model_metadata;
|
||||
pub mod models;
|
||||
pub mod utils;
|
||||
|
||||
use anyhow::{Result, anyhow};
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use image::DynamicImage;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
pub use crate::models::det::{Detector,DetectionResult};
|
||||
pub use crate::models::ocr::{Ocr, OcrPredictor, OcrResult};
|
||||
pub use crate::models::slide::{Slider, SlideResult};
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
// 关键点：直接使用 tract 重导出的 ndarray
|
||||
use crate::charset::CharRestrict;
|
||||
use crate::model_metadata::ModelMetadata;
|
||||
use crate::models::det::DetectionResult;
|
||||
pub use crate::model_metadata::ModelMetadata;
|
||||
use crate::utils::color_filter::{ColorPreset, HsvRange};
|
||||
use models::det::Det;
|
||||
use models::loader::ModelSession;
|
||||
use models::ocr::Ocr;
|
||||
|
||||
pub enum ModelSpec {
|
||||
/// 默认 OCR (使用内置路径)
|
||||
OcrModel,
|
||||
DetModel,
|
||||
/// 自定义 OCR (路径由用户提供)
|
||||
CustomOcrModel {
|
||||
path: String,
|
||||
model_metadata: ModelMetadata,
|
||||
},
|
||||
}
|
||||
impl ModelSpec {
|
||||
// 将默认路径定义为内部关联常量
|
||||
const DEFAULT_OCR_PATH: &'static str = "models/common_sml2h3_f32.onnx";
|
||||
const DEFAULT_DET_PATH: &'static str = "models/common_det.onnx";
|
||||
}
|
||||
pub enum Runtime {
|
||||
Ocr(Ocr),
|
||||
Det(Det),
|
||||
}
|
||||
impl Runtime {
|
||||
// 统一获取描述的方法
|
||||
pub fn desc(&self) -> String {
|
||||
match self {
|
||||
Runtime::Ocr(s) => s.desc(), // 调用 Ocr 结构体的方法
|
||||
Runtime::Det(s) => s.desc(), // 调用 Det 结构体的方法
|
||||
}
|
||||
}
|
||||
}
|
||||
pub struct DdddOcrBuilder {
|
||||
mode: ModelSpec,
|
||||
}
|
||||
|
||||
impl DdddOcrBuilder {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
mode: ModelSpec::OcrModel,
|
||||
}
|
||||
}
|
||||
|
||||
/// 切换为检测模式
|
||||
pub fn det(mut self) -> Self {
|
||||
self.mode = ModelSpec::DetModel;
|
||||
self
|
||||
}
|
||||
|
||||
/// 设置自定义 OCR 路径
|
||||
pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self {
|
||||
// 直接重写枚举，替换掉之前的 Ocr 或 Det
|
||||
self.mode = ModelSpec::CustomOcrModel {
|
||||
path,
|
||||
model_metadata,
|
||||
};
|
||||
self
|
||||
}
|
||||
|
||||
/// 核心初始化逻辑
|
||||
pub fn build(self) -> Result<DdddOcr> {
|
||||
let runtime = match self.mode {
|
||||
ModelSpec::OcrModel => Runtime::Ocr(Ocr::new(
|
||||
ModelSpec::DEFAULT_OCR_PATH.into(),
|
||||
ModelMetadata::from_builtin_beta(),
|
||||
)?),
|
||||
ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?),
|
||||
ModelSpec::CustomOcrModel {
|
||||
path,
|
||||
model_metadata,
|
||||
} => Runtime::Ocr(Ocr::new(path, model_metadata)?),
|
||||
};
|
||||
|
||||
Ok(DdddOcr { runtime })
|
||||
}
|
||||
}
|
||||
|
||||
pub struct DdddOcr {
|
||||
runtime: Runtime,
|
||||
}
|
||||
|
||||
impl Display for DdddOcr {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
write!(f, "DdddOcr(session: {})", self.runtime.desc())
|
||||
}
|
||||
}
|
||||
|
||||
impl DdddOcr {
|
||||
pub fn classification(&self, img: &DynamicImage) -> Result<String> {
|
||||
match &self.runtime {
|
||||
// Runtime::Ocr(s) => s.predict(img).run(),
|
||||
// Runtime::Ocr(s) => s.predictor().probability(false).predict(img),
|
||||
// Runtime::Ocr(s) => {
|
||||
// let predictor = s.predictor();
|
||||
// let restricted = predictor.charset_restrict(&CharRestrict::Lowercase);
|
||||
// let a = restricted.valid_tokens();
|
||||
// println!("{:?}", a);
|
||||
// Ok("".to_string())
|
||||
// }
|
||||
Runtime::Ocr(s) => {
|
||||
let res = s.predictor().probability(true).predict(img)?;
|
||||
println!("{}", res);
|
||||
Ok(res.to_string())
|
||||
}
|
||||
// Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img),
|
||||
// Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![
|
||||
// // 错误：下界 (82, 221, 14) 没问题
|
||||
// // 但上界的 H 通道写成了 240，超过了 180 的法定上限！
|
||||
// HsvRange::new((82, 221, 14), (240, 203, 82)),
|
||||
// ])).predict(img),
|
||||
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")),
|
||||
}
|
||||
}
|
||||
pub fn detection(&self, img: &DynamicImage) -> Result<Vec<DetectionResult>> {
|
||||
match &self.runtime {
|
||||
Runtime::Det(s) => s.predict(img),
|
||||
Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")),
|
||||
}
|
||||
}
|
||||
}
|
||||
// pub enum ModelSpec {
|
||||
// /// 默认 OCR (使用内置路径)
|
||||
// OcrModel,
|
||||
// DetModel,
|
||||
// /// 自定义 OCR (路径由用户提供)
|
||||
// CustomOcrModel {
|
||||
// path: String,
|
||||
// model_metadata: ModelMetadata,
|
||||
// },
|
||||
// }
|
||||
// impl ModelSpec {
|
||||
// // 将默认路径定义为内部关联常量
|
||||
// const DEFAULT_OCR_PATH: &'static str = "models/common_sml2h3_f32.onnx";
|
||||
// const DEFAULT_DET_PATH: &'static str = "models/common_det.onnx";
|
||||
// }
|
||||
// pub enum Runtime {
|
||||
// Ocr(Ocr),
|
||||
// Det(Det),
|
||||
// }
|
||||
// impl Runtime {
|
||||
// // 统一获取描述的方法
|
||||
// pub fn desc(&self) -> String {
|
||||
// match self {
|
||||
// Runtime::Ocr(s) => s.desc(), // 调用 Ocr 结构体的方法
|
||||
// Runtime::Det(s) => s.desc(), // 调用 Det 结构体的方法
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
// pub struct DdddOcrBuilder {
|
||||
// mode: ModelSpec,
|
||||
// }
|
||||
//
|
||||
// impl DdddOcrBuilder {
|
||||
// pub fn new() -> Self {
|
||||
// Self {
|
||||
// mode: ModelSpec::OcrModel,
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// /// 切换为检测模式
|
||||
// pub fn det(mut self) -> Self {
|
||||
// self.mode = ModelSpec::DetModel;
|
||||
// self
|
||||
// }
|
||||
//
|
||||
// /// 设置自定义 OCR 路径
|
||||
// pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self {
|
||||
// // 直接重写枚举，替换掉之前的 Ocr 或 Det
|
||||
// self.mode = ModelSpec::CustomOcrModel {
|
||||
// path,
|
||||
// model_metadata,
|
||||
// };
|
||||
// self
|
||||
// }
|
||||
//
|
||||
// /// 核心初始化逻辑
|
||||
// pub fn build(self) -> Result<DdddOcr> {
|
||||
// let runtime = match self.mode {
|
||||
// ModelSpec::OcrModel => Runtime::Ocr(Ocr::new(
|
||||
// ModelSpec::DEFAULT_OCR_PATH.into(),
|
||||
// ModelMetadata::from_builtin_beta(),
|
||||
// )?),
|
||||
// ModelSpec::DetModel => Runtime::Det(Det::new(ModelSpec::DEFAULT_DET_PATH.into())?),
|
||||
// ModelSpec::CustomOcrModel {
|
||||
// path,
|
||||
// model_metadata,
|
||||
// } => Runtime::Ocr(Ocr::new(path, model_metadata)?),
|
||||
// };
|
||||
//
|
||||
// Ok(DdddOcr { runtime })
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// pub struct DdddOcr {
|
||||
// runtime: Runtime,
|
||||
// }
|
||||
//
|
||||
// impl Display for DdddOcr {
|
||||
// fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
// write!(f, "DdddOcr(session: {})", self.runtime.desc())
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// impl DdddOcr {
|
||||
// pub fn classification(&self, img: &DynamicImage) -> Result<String> {
|
||||
// match &self.runtime {
|
||||
// // Runtime::Ocr(s) => s.predict(img).run(),
|
||||
// // Runtime::Ocr(s) => s.predictor().probability(false).predict(img),
|
||||
// // Runtime::Ocr(s) => {
|
||||
// // let predictor = s.predictor();
|
||||
// // let restricted = predictor.charset_restrict(&CharRestrict::Lowercase);
|
||||
// // let a = restricted.valid_tokens();
|
||||
// // println!("{:?}", a);
|
||||
// // Ok("".to_string())
|
||||
// // }
|
||||
// Runtime::Ocr(s) => {
|
||||
// let res = s.predictor().probability(true).predict(img)?;
|
||||
// println!("{}", res);
|
||||
// Ok(res.to_string())
|
||||
// }
|
||||
// // Runtime::Ocr(s) => s.predictor().charset_restrict(&CharRestrict::Digit).predict(img),
|
||||
// // Runtime::Ocr(s) => s.predictor().color_filter(&ColorPreset::Custom(vec![
|
||||
// // // 错误：下界 (82, 221, 14) 没问题
|
||||
// // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！
|
||||
// // HsvRange::new((82, 221, 14), (240, 203, 82)),
|
||||
// // ])).predict(img),
|
||||
// Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")),
|
||||
// }
|
||||
// }
|
||||
// pub fn detection(&self, img: &DynamicImage) -> Result<Vec<DetectionResult>> {
|
||||
// match &self.runtime {
|
||||
// Runtime::Det(s) => s.predict(img),
|
||||
// Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型，无法执行检测")),
|
||||
// }
|
||||
// }
|
||||
// }
|
||||
|
||||
// struct Classification {}
|
||||
// #[derive(Debug)]
|
||||
|
||||
Reference in New Issue
Block a user