refactor(ocr): 优化 color_filter.rs

- 重构 `OcrBuilder` 移除is_pixel_matched,filter_image。
 - 优化 `OcrBuilder` 的color_filter方法(部分逻辑转移给merge_to_vec) 。
 - 新增 `ColorFilter` 特征增加merge_to_vec方法。
This commit is contained in:
2026-06-25 20:25:49 +08:00
parent 62d5e7a0ca
commit 2f86694c54
4 changed files with 184 additions and 67 deletions

View File

@@ -1,18 +1,21 @@
use std::borrow::Cow;
use crate::charset::{TokenFilter, ValidationCtx};
use crate::model_metadata::ModelMetadata;
use crate::models::base::ModelArgs;
use crate::models::loader::{ModelLoader, ModelSession, ModelType};
use crate::utils::color_filter::{ColorFilter, HsvRange};
use crate::utils::color_filter::{filter_image, ColorFilter, HsvRange};
use crate::utils::image_io::png_rgba_white_preprocess;
use crate::utils::image_processor::{convert_to_grayscale, resize_image};
use anyhow::Context;
use image::DynamicImage;
use anyhow::{anyhow, Result};
use image::{DynamicImage, ImageBuffer, Rgb};
use std::collections::HashSet;
use tract_onnx::prelude::tract_ndarray::s;
use tract_onnx::prelude::{
DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec,
};
// 引入 cv_ops 模块中的 OpenCV HSV 转换算子
use crate::utils::cv_ops::rgb_to_opencv_hsv;
// 颜色过滤的自定义范围:(低值RGB, 高值RGB)
pub type ColorRange = ((u8, u8, u8), (u8, u8, u8));
@@ -185,19 +188,16 @@ impl<'a> OcrBuilder<'a> {
existing_vec.reserve(total_capacity);
// 尝试追加倒入
filter.append_ranges(existing_vec) ;
filter.append_ranges(existing_vec);
}
}
if let Some(v) =&mut ranges {
if let Some(v) = &mut ranges {
v.sort_unstable();
v.dedup();
}
// 5. 更新状态
self.color_filter = Ok(ranges);
},
}
Err(_) => return self,
};
self
@@ -252,7 +252,27 @@ impl<'a> OcrBuilder<'a> {
}
impl<'a> OcrBuilder<'a> {
pub fn predict(&self, image: &DynamicImage) -> anyhow::Result<String> {
let tensor = self.preprocess_image(image)?;
println!("当前颜色过滤器状态: {:?}", self.color_filter);
// =====================================================================
// 管道节点 1: 颜色过滤流水线
// 使用 Cow (Copy-On-Write) 智能指针。
// 如果未开启过滤img_cow 内部只是持有原图的【只读借用】,发生【零内存分配】!
// =====================================================================
let img_cow = match &self.color_filter {
Err(err_msg) => {
return Err(anyhow::anyhow!("颜色过滤器初始化失败,全链路短路: {}", err_msg));
}
Ok(None) => {
// 核心优化点:直接借用原图,不发生任何克隆
Cow::Borrowed(image)
}
Ok(Some(ranges)) => {
// 只有真正需要过滤时,才在内部提取像素并生成清洗后的 Owned 新图
let filtered_img = filter_image(image, ranges)?;
Cow::Owned(filtered_img)
}
};
let tensor = self.preprocess_image(&img_cow)?;
let raw_tensor = self.ocr.inference(tensor)?;
let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?;
@@ -266,15 +286,17 @@ impl<'a> OcrBuilder<'a> {
/// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换
fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result<Tensor> {
// A. 修复 PNG 透明背景 (内部逻辑你之前已实现)
let _ = if self.png_fix && img.color().has_alpha() {
png_rgba_white_preprocess(img)
let current_img = if self.png_fix && img.color().has_alpha() {
// 只有满足条件才去触发分配,生成新图
Cow::Owned(png_rgba_white_preprocess(img))
} else {
img.clone()
// 正常情况下,仅仅是再次安全借用,无开销
Cow::Borrowed(img)
};
let h = 64u32;
let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32;
let gray_img = convert_to_grayscale(img);
let w = (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32;
let gray_img = convert_to_grayscale(&current_img);
let resized = resize_image(&gray_img, w, h);
// resized.save("debug_preprocessed.png").unwrap();
// 1. 预处理:转灰度 -> Resize -> 归一化
@@ -291,6 +313,10 @@ impl<'a> OcrBuilder<'a> {
Ok(tensor)
}
}
impl<'a> OcrBuilder<'a> {
pub fn get_valid_indices(&self) -> HashSet<usize> {