188 lines
6.3 KiB
Rust
188 lines
6.3 KiB
Rust
mod charset;
|
||
mod error;
|
||
mod model_metadata;
|
||
pub mod models;
|
||
pub mod utils;
|
||
|
||
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;
|
||
pub use crate::model_metadata::ModelMetadata;
|
||
use crate::utils::color_filter::{ColorPreset, HsvRange};
|
||
use models::loader::ModelSession;
|
||
|
||
|
||
// 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)]
|
||
// struct ClassificationBuilder {
|
||
// img: DynamicImage,
|
||
// png_fix: bool,
|
||
// color_filter_colors: Option<Vec<ColorRange>>,
|
||
// color_filter_custom_ranges: Option<Vec<ColorRange>>,
|
||
// }
|
||
// impl ClassificationBuilder {
|
||
// pub fn new(img: DynamicImage) -> Self {
|
||
// ClassificationBuilder {
|
||
// img,
|
||
// png_fix: false,
|
||
// color_filter_colors: None,
|
||
// color_filter_custom_ranges: None,
|
||
// }
|
||
// }
|
||
// pub fn png_fix(mut self, value: bool) -> Self {
|
||
// self.png_fix = value;
|
||
// self
|
||
// }
|
||
// pub fn color_filter_colors(mut self, value: Vec<ColorRange>) -> Self {
|
||
// self.color_filter_colors = Some(value);
|
||
// self
|
||
// }
|
||
// pub fn color_filter_custom_ranges(mut self, value: Vec<ColorRange>) -> Self {
|
||
// self.color_filter_custom_ranges = Some(value);
|
||
// self
|
||
// }
|
||
// pub fn build(self) -> Classification {
|
||
// Classification {}
|
||
// }
|
||
// }
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
#[test]
|
||
fn test_ctc_decode_indices() {
|
||
// 模拟一个 DdddOcr 实例（如果 decode 不依赖 session，可以设为相关函数）
|
||
// 这里假设你的 decode_ctc 是公开或内部可访问的
|
||
let input = vec![1, 1, 0, 1, 2, 2, 0, 2];
|
||
// 逻辑：[1, 1] -> 1, [0] -> 跳过, [1] -> 1, [2, 2] -> 2, [0] -> 跳过, [2] -> 2
|
||
// 预期结果索引应该是 [1, 1, 2, 2] 对应的字符
|
||
// 具体的断言取决于你的 CHARSET_BETA
|
||
// let result = dddd.ctc_decode_indices(&input);
|
||
// assert_eq!(result, "AABB");
|
||
}
|
||
}
|