refactor(predict): 重构预测流水线并优化模型元数据与输出架构
- 优化 `predict` 核心方法:移除冗余日志与深层嵌套,将流程重塑为线性流水线。 - 重构 `compute_f32_full_probability`:解耦逻辑与外部状态,消除并发隐患与生命周期冲突。 - 增强 `ModelMetadata`:引入动态归一化配置并支持 Serde 序列化,解决特定模型漏字问题。 - 升级 `OcrOutput`: - 增加 `Unsupported` 变体以支持非致命异常的优雅降级。 - 实现 `into_text(self)` 方法与 `Display` 特征(应用双重截断保护,防止日志刷屏)。 BREAKING CHANGE: `predict` 返回值由 `anyhow::Result<String>` 改为 `anyhow::Result<OcrOutput>`,将后处理和控制权移交上层。
This commit is contained in:
49
src/lib.rs
49
src/lib.rs
@@ -1,20 +1,20 @@
|
||||
mod charset;
|
||||
|
||||
mod model_metadata;
|
||||
pub mod models;
|
||||
pub mod utils;
|
||||
mod model_metadata;
|
||||
|
||||
use anyhow::Result;
|
||||
use anyhow::{Result, anyhow};
|
||||
use image::DynamicImage;
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
// 关键点:直接使用 tract 重导出的 ndarray
|
||||
use crate::charset::{ CharRestrict};
|
||||
use crate::charset::CharRestrict;
|
||||
use crate::model_metadata::ModelMetadata;
|
||||
use crate::utils::color_filter::{ColorPreset, HsvRange};
|
||||
use models::det::Det;
|
||||
use models::loader::ModelSession;
|
||||
use models::ocr::Ocr;
|
||||
use crate::model_metadata::ModelMetadata;
|
||||
use crate::utils::color_filter::{ColorPreset, HsvRange};
|
||||
|
||||
pub enum ModelSpec {
|
||||
/// 默认 OCR (使用内置路径)
|
||||
@@ -64,7 +64,10 @@ impl DdddOcrBuilder {
|
||||
/// 设置自定义 OCR 路径
|
||||
pub fn custom_ocr(mut self, path: String, model_metadata: ModelMetadata) -> Self {
|
||||
// 直接重写枚举,替换掉之前的 Ocr 或 Det
|
||||
self.mode = ModelSpec::CustomOcrModel { path, model_metadata };
|
||||
self.mode = ModelSpec::CustomOcrModel {
|
||||
path,
|
||||
model_metadata,
|
||||
};
|
||||
self
|
||||
}
|
||||
|
||||
@@ -76,7 +79,10 @@ impl DdddOcrBuilder {
|
||||
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)?),
|
||||
ModelSpec::CustomOcrModel {
|
||||
path,
|
||||
model_metadata,
|
||||
} => Runtime::Ocr(Ocr::new(path, model_metadata)?),
|
||||
};
|
||||
|
||||
Ok(DdddOcr { runtime })
|
||||
@@ -97,23 +103,36 @@ 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().predict(img),
|
||||
// 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("".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")),
|
||||
}
|
||||
|
||||
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型,无法执行 OCR")),
|
||||
}
|
||||
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> {
|
||||
match &self.runtime {
|
||||
Runtime::Det(s) => s.predict(img),
|
||||
Runtime::Ocr(_) => Err(anyhow::anyhow!("当前模型是 OCR 模型,无法执行检测")),
|
||||
}
|
||||
}
|
||||
pub fn detection(&self, img: &[u8]) -> Result<Vec<Vec<i32>>> {
|
||||
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