refactor: 重构 ocr.rs 实现丝滑参数设置

- 重构 Ocr
- 新增 OcrTask

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<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]

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<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
-    charset: Vec<String>,
+    pub session: RunnableModel<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
+    pub charset: Vec<String>,
 }
 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<Vec<ColorRange>>,
+    color_filter_custom_ranges: Option<Vec<ColorRange>>,
+}
+
+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<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 predict(&self, image: &DynamicImage, png_fix: bool) -> Result<String, anyhow::Error> {
         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<Tensor> {
@@ -156,6 +194,7 @@ impl Ocr {
         // tract 的 run 会返回一个 Vec<TValue>，我们通常只需要第一个输出
         // 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 {
                         // 保护逻辑：如果模型预测的索引超出了字符集范围