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

- 重构 Ocr
- 新增 OcrTask

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 {
                         // 保护逻辑：如果模型预测的索引超出了字符集范围