diff --git a/src/lib.rs b/src/lib.rs index 391554a..8899498 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -9,9 +9,11 @@ use std::fmt::{Display, Formatter}; // 关键点:直接使用 tract 重导出的 ndarray use crate::charset::get_default_charset; +use crate::models::ocr::ColorRange; use models::det::Det; use models::loader::ModelSession; use models::ocr::Ocr; + pub enum ModelSpec { /// 默认 OCR (使用内置路径) OcrModel, @@ -24,7 +26,7 @@ pub enum ModelSpec { } impl ModelSpec { // 将默认路径定义为内部关联常量 - const DEFAULT_OCR_PATH: &'static str = "models/common.onnx"; + const DEFAULT_OCR_PATH: &'static str = "models/common_sml2h3_f32.onnx"; const DEFAULT_DET_PATH: &'static str = "models/common_det.onnx"; } pub enum Runtime { @@ -104,6 +106,40 @@ impl DdddOcr { } } +struct Classification {} +#[derive(Debug)] +struct ClassificationBuilder { + img: DynamicImage, + png_fix: bool, + color_filter_colors: Option>, + color_filter_custom_ranges: Option>, +} +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) -> Self { + self.color_filter_colors = Some(value); + self + } + pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { + self.color_filter_custom_ranges = Some(value); + self + } + pub fn build(self) -> Classification { + Classification {} + } +} + #[cfg(test)] mod tests { #[test] diff --git a/src/models/ocr.rs b/src/models/ocr.rs index e9c39d2..e287b5f 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -1,7 +1,7 @@ use crate::models::base::ModelArgs; +use crate::models::loader::{ModelLoader, ModelSession, ModelType}; use crate::utils::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; -use crate::models::loader::{ModelLoader, ModelSession, ModelType}; use anyhow::Context; use image::DynamicImage; use tract_onnx::prelude::tract_ndarray::s; @@ -95,8 +95,8 @@ impl PredictArgs { } } pub struct Ocr { - session: RunnableModel, Graph>>, - charset: Vec, + pub session: RunnableModel, Graph>>, + pub charset: Vec, } impl ModelSession for Ocr { fn get_model_type(&self) -> ModelType { @@ -111,6 +111,43 @@ impl Ocr { let session = ModelLoader::load_model(&model_path)?.session; Ok(Self { session, charset }) } + pub fn task<'a>(&'a self, image: &'a DynamicImage) -> OcrTask { + OcrTask::new(self, image) + } +} + +pub struct OcrTask<'a> { + ocr: &'a Ocr, + image: &'a DynamicImage, + png_fix: bool, + color_filter_colors: Option>, + color_filter_custom_ranges: Option>, +} + +impl<'a> OcrTask<'a> { + // 初始化任务,设置默认参数 + pub fn new(ocr: &'a Ocr, image: &'a DynamicImage) -> Self { + Self { + ocr, + image, + 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) -> Self { + self.color_filter_colors = Some(value); + self + } + pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { + self.color_filter_custom_ranges = Some(value); + self + } + pub fn predict(&self, image: &DynamicImage, png_fix: bool) -> Result { let tensor = self.preprocess_image(image, png_fix)?; // @@ -122,6 +159,7 @@ impl Ocr { Ok(self.ctc_decode_indices(&output2)) // Ok("ocr result".to_string()) } + /// 对应 Python 的 _preprocess_image /// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 fn preprocess_image(&self, img: &DynamicImage, png_fix: bool) -> anyhow::Result { @@ -156,6 +194,7 @@ impl Ocr { // tract 的 run 会返回一个 Vec,我们通常只需要第一个输出 // let result = self.session.run(tvec!(tensor.into()))?; let mut result = self + .ocr .session .run(tvec!(tensor.into())) .context("执行模型推理失败")?; @@ -235,7 +274,7 @@ impl Ocr { // 2. 跳过 blank 字符 (假设索引 0 是 blank) if idx != prev_idx && idx != 0 { if let Ok(u_idx) = usize::try_from(idx) { - if let Some(char_str) = self.charset.get(u_idx) { + if let Some(char_str) = self.ocr.charset.get(u_idx) { res.push_str(char_str); } else { // 保护逻辑:如果模型预测的索引超出了字符集范围