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

@@ -98,12 +98,13 @@ impl DdddOcr {
pub fn classification(&self, img: &DynamicImage) -> Result<String> { pub fn classification(&self, img: &DynamicImage) -> Result<String> {
match &self.runtime { match &self.runtime {
// Runtime::Ocr(s) => s.predict(img).run(), // Runtime::Ocr(s) => s.predict(img).run(),
Runtime::Ocr(s) => s.builder().predict(img),
// Runtime::Ocr(s) => s.builder().charset_restrict(&CharRestrict::Digit).predict(img), // Runtime::Ocr(s) => s.builder().charset_restrict(&CharRestrict::Digit).predict(img),
Runtime::Ocr(s) => s.builder().color_filter(&ColorPreset::Custom(vec![ // Runtime::Ocr(s) => s.builder().color_filter(&ColorPreset::Custom(vec![
// 错误:下界 (82, 221, 14) 没问题 // // 错误:下界 (82, 221, 14) 没问题
// 但上界的 H 通道写成了 240超过了 180 的法定上限! // // 但上界的 H 通道写成了 240超过了 180 的法定上限!
HsvRange::new((82, 221, 14), (240, 203, 82)), // HsvRange::new((82, 221, 14), (240, 203, 82)),
])).predict(img), // ])).predict(img),
Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型,无法执行 OCR")), Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型,无法执行 OCR")),
} }
} }

View File

@@ -1,18 +1,21 @@
use std::borrow::Cow;
use crate::charset::{TokenFilter, ValidationCtx}; use crate::charset::{TokenFilter, ValidationCtx};
use crate::model_metadata::ModelMetadata; use crate::model_metadata::ModelMetadata;
use crate::models::base::ModelArgs; use crate::models::base::ModelArgs;
use crate::models::loader::{ModelLoader, ModelSession, ModelType}; 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_io::png_rgba_white_preprocess;
use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use crate::utils::image_processor::{convert_to_grayscale, resize_image};
use anyhow::Context; use anyhow::Context;
use image::DynamicImage; use anyhow::{anyhow, Result};
use image::{DynamicImage, ImageBuffer, Rgb};
use std::collections::HashSet; use std::collections::HashSet;
use tract_onnx::prelude::tract_ndarray::s; use tract_onnx::prelude::tract_ndarray::s;
use tract_onnx::prelude::{ use tract_onnx::prelude::{
DatumType, Graph, IntoTensor, RunnableModel, Tensor, TypedFact, TypedOp, tract_ndarray, tvec, 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) // 颜色过滤的自定义范围:(低值RGB, 高值RGB)
pub type ColorRange = ((u8, u8, u8), (u8, u8, u8)); pub type ColorRange = ((u8, u8, u8), (u8, u8, u8));
@@ -186,10 +189,7 @@ impl<'a> OcrBuilder<'a> {
// 尝试追加倒入 // 尝试追加倒入
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.sort_unstable();
@@ -197,7 +197,7 @@ impl<'a> OcrBuilder<'a> {
} }
// 5. 更新状态 // 5. 更新状态
self.color_filter = Ok(ranges); self.color_filter = Ok(ranges);
}, }
Err(_) => return self, Err(_) => return self,
}; };
self self
@@ -252,7 +252,27 @@ impl<'a> OcrBuilder<'a> {
} }
impl<'a> OcrBuilder<'a> { impl<'a> OcrBuilder<'a> {
pub fn predict(&self, image: &DynamicImage) -> anyhow::Result<String> { 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_tensor = self.ocr.inference(tensor)?;
let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?; let raw_indices = self.ocr.extract_indices_from_tensor(&raw_tensor)?;
@@ -266,15 +286,17 @@ impl<'a> OcrBuilder<'a> {
/// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换 /// 负责:透明背景修复 -> 灰度化 -> 按比例 Resize -> 归一化 -> 4维张量转换
fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result<Tensor> { fn preprocess_image(&self, img: &DynamicImage) -> anyhow::Result<Tensor> {
// A. 修复 PNG 透明背景 (内部逻辑你之前已实现) // A. 修复 PNG 透明背景 (内部逻辑你之前已实现)
let _ = if self.png_fix && img.color().has_alpha() { let current_img = if self.png_fix && img.color().has_alpha() {
png_rgba_white_preprocess(img) // 只有满足条件才去触发分配,生成新图
Cow::Owned(png_rgba_white_preprocess(img))
} else { } else {
img.clone() // 正常情况下,仅仅是再次安全借用,无开销
Cow::Borrowed(img)
}; };
let h = 64u32; let h = 64u32;
let w = (img.width() as f32 * (h as f32 / img.height() as f32)) as u32; let w = (current_img.width() as f32 * (h as f32 / current_img.height() as f32)) as u32;
let gray_img = convert_to_grayscale(img); let gray_img = convert_to_grayscale(&current_img);
let resized = resize_image(&gray_img, w, h); let resized = resize_image(&gray_img, w, h);
// resized.save("debug_preprocessed.png").unwrap(); // resized.save("debug_preprocessed.png").unwrap();
// 1. 预处理:转灰度 -> Resize -> 归一化 // 1. 预处理:转灰度 -> Resize -> 归一化
@@ -291,6 +313,10 @@ impl<'a> OcrBuilder<'a> {
Ok(tensor) Ok(tensor)
} }
} }
impl<'a> OcrBuilder<'a> { impl<'a> OcrBuilder<'a> {
pub fn get_valid_indices(&self) -> HashSet<usize> { pub fn get_valid_indices(&self) -> HashSet<usize> {

View File

@@ -1,4 +1,7 @@
use std::str::FromStr; use std::str::FromStr;
use anyhow::anyhow;
use image::{DynamicImage, ImageBuffer, Rgb};
use crate::utils::cv_ops::rgb_to_opencv_hsv;
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct HsvRange { pub struct HsvRange {
@@ -6,6 +9,47 @@ pub struct HsvRange {
pub upper: (u8, u8, u8), // (H, S, V) pub upper: (u8, u8, u8), // (H, S, V)
} }
/// 核心区间判定辅助函数
#[inline(always)]
fn is_pixel_matched(ranges: &[HsvRange], h: u8, s: u8, v: u8) -> bool {
ranges.iter().any(|range| {
h >= range.lower.0 && h <= range.upper.0 &&
s >= range.lower.1 && s <= range.upper.1 &&
v >= range.lower.2 && v <= range.upper.2
})
}
pub fn filter_image(image: &DynamicImage, hsv_ranges: &[HsvRange]) -> anyhow::Result<DynamicImage> {
// 1. 统一转换为连续内存的 RGB8 缓冲区 (对应 Python 的 Image 到 RGB/BGR 数组转换)
let rgb_img = image.to_rgb8();
let (width, height) = rgb_img.dimensions();
let mut raw_pixels = rgb_img.into_raw();
// 2. 密集计算核心:原地流式迭代修改
// 每次取出 3 个 u8 字节,分别代表 [R, G, B],无多余掩膜矩阵内存分配
for chunk in raw_pixels.chunks_exact_mut(3) {
let r = chunk[0];
let g = chunk[1];
let b = chunk[2];
// 像素级转换为 OpenCV 标准的 HSV
let (h, s, v) = rgb_to_opencv_hsv(r, g, b);
// 模拟 Python 的多范围 mask bitwise_or 并在 mask == 0 处刷白
// 如果该像素没有命中任何一个配置的颜色区间,立刻原地刷白 [255, 255, 255]
if !is_pixel_matched(hsv_ranges, h, s, v) {
chunk[0] = 255;
chunk[1] = 255;
chunk[2] = 255;
}
}
// 3. 将扁平字节数组重新打包回 DynamicImage 容器
let filtered_buffer = ImageBuffer::<Rgb<u8>, Vec<u8>>::from_raw(width, height, raw_pixels)
.ok_or_else(|| anyhow!("图像缓冲重新组装失败,维度与数据大小不匹配"))?;
Ok(DynamicImage::ImageRgb8(filtered_buffer))
}
impl HsvRange { impl HsvRange {
pub const fn new(lower: (u8, u8, u8), upper: (u8, u8, u8)) -> Self { pub const fn new(lower: (u8, u8, u8), upper: (u8, u8, u8)) -> Self {
Self { lower, upper } Self { lower, upper }
@@ -119,6 +163,38 @@ pub trait ColorFilter {
fn validate_self(&self) -> Result<(), String> { fn validate_self(&self) -> Result<(), String> {
Ok(()) Ok(())
} }
/// 【新扩展的架构方法】将自身安全的合并到已有的普通容器中,并完成去重和排序
/// 完美的责任分离Builder 不再需要关心怎么分配内存、怎么排序去重
fn merge_to_vec(&self, mut existing: Option<Vec<HsvRange>>) -> Result<Option<Vec<HsvRange>>, String> {
// 1. 触发自检
self.validate_self()?;
let total_capacity = self.estimated_count();
if total_capacity == 0 {
return Ok(existing);
}
let mut v = match existing {
None => {
// 情况 A第一次配置精准分配
Vec::with_capacity(total_capacity)
}
Some(mut existing_vec) => {
// 情况 B追加配置精准扩容
existing_vec.reserve(total_capacity);
existing_vec
}
};
// 2. 倒入数据
self.append_ranges(&mut v);
// 3. 原地完成排序与去重
v.sort_unstable();
v.dedup();
Ok(Some(v))
}
} }
impl ColorFilter for ColorPreset { impl ColorFilter for ColorPreset {
@@ -180,43 +256,3 @@ macro_rules! color_any_of {
}; };
} }
// =====================================================================
// 5. 核心高性能图像转换算法 (纯 Rust 编写)
// =====================================================================
/// 极速无拷贝 RGB 转 HSV 算法 (完全对齐 OpenCV 行为)
#[inline]
pub fn rgb_to_hsv(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
let r_f = r as f32 / 255.0;
let g_f = g as f32 / 255.0;
let b_f = b as f32 / 255.0;
let max = r_f.max(g_f).max(b_f);
let min = r_f.min(g_f).min(b_f);
let delta = max - min;
// 1. 计算 H (色调)
let mut h = if delta == 0.0 {
0.0
} else if max == r_f {
60.0 * (((g_f - b_f) / delta) % 6.0)
} else if max == g_f {
60.0 * (((b_f - r_f) / delta) + 2.0)
} else {
60.0 * (((r_f - g_f) / delta) + 4.0)
};
if h < 0.0 {
h += 360.0;
}
let h_opencv = (h / 2.0).round() as u8;
// 2. 计算 S (饱和度)
let s = if max == 0.0 { 0.0 } else { delta / max };
let s_opencv = (s * 255.0).round() as u8;
// 3. 计算 V (明度)
let v_opencv = (max * 255.0).round() as u8;
(h_opencv, s_opencv, v_opencv)
}

View File

@@ -1,6 +1,6 @@
use std::cmp::{max, min};
use image::{ImageBuffer, Luma}; use image::{ImageBuffer, Luma};
use tract_onnx::prelude::tract_ndarray::{azip, Array2, Array3, ArrayView2, ArrayView3}; use std::cmp::{max, min};
use tract_onnx::prelude::tract_ndarray::{Array2, Array3, ArrayView2, ArrayView3, azip};
/// 1. 计算两个数组的绝对差值 (对应 cv2.absdiff) /// 1. 计算两个数组的绝对差值 (对应 cv2.absdiff)
pub fn abs_diff(a: &ArrayView3<u8>, b: &ArrayView3<u8>) -> Array3<u8> { pub fn abs_diff(a: &ArrayView3<u8>, b: &ArrayView3<u8>) -> Array3<u8> {
@@ -13,7 +13,6 @@ pub fn abs_diff(a: &ArrayView3<u8>, b: &ArrayView3<u8>) -> Array3<u8> {
diff diff
} }
/// RGB 到灰度转换 /// RGB 到灰度转换
pub fn rgb_to_gray(rgb: ArrayView3<u8>) -> Array2<u8> { pub fn rgb_to_gray(rgb: ArrayView3<u8>) -> Array2<u8> {
let (h, w, _) = rgb.dim(); let (h, w, _) = rgb.dim();
@@ -67,9 +66,16 @@ pub fn find_contours_and_max(labelled: &ImageBuffer<Luma<u32>, Vec<u32>>) -> Opt
max_label = label; max_label = label;
} }
} }
if max_label == 0 { None } else { Some(max_label) } if max_label == 0 {
None
} else {
Some(max_label)
} }
pub fn bounding_rect(labelled: &ImageBuffer<Luma<u32>, Vec<u32>>,max_label: u32) -> (u32, u32, u32, u32) { }
pub fn bounding_rect(
labelled: &ImageBuffer<Luma<u32>, Vec<u32>>,
max_label: u32,
) -> (u32, u32, u32, u32) {
// 5. 计算最大区域的边界框 (对应 cv2.boundingRect) // 5. 计算最大区域的边界框 (对应 cv2.boundingRect)
let mut min_x = labelled.width(); let mut min_x = labelled.width();
let mut max_x = 0; let mut max_x = 0;
@@ -85,7 +91,6 @@ pub fn bounding_rect(labelled: &ImageBuffer<Luma<u32>, Vec<u32>>,max_label: u32)
} }
} }
let w = max_x - min_x; let w = max_x - min_x;
let h = max_y - min_y; let h = max_y - min_y;
(min_x, min_y, w, h) (min_x, min_y, w, h)
@@ -105,3 +110,52 @@ pub fn ndarray_to_luma8(array: ArrayView2<u8>) -> ImageBuffer<Luma<u8>, Vec<u8>>
} }
buffer buffer
} }
// =====================================================================
// 5. 核心高性能图像转换算法 (纯 Rust 编写)
// =====================================================================
#[inline(always)]
pub fn rgb_to_opencv_hsv(r: u8, g: u8, b: u8) -> (u8, u8, u8) {
// 1. 规避高昂的除法,直接转为 f32 进行比对
let r_f = r as f32;
let g_f = g as f32;
let b_f = b as f32;
let max = r_f.max(g_f).max(b_f);
let min = r_f.min(g_f).min(b_f);
let delta = max - min;
// 2. 计算 H (色调) - 移除负数取余陷阱,改用平铺分支
let mut h = if delta == 0.0 {
0.0
} else if max == r_f {
let mut diff = (g_f - b_f) / delta;
if diff < 0.0 {
diff += 6.0; // 规避 Rust f32 % 负数的行为
}
60.0 * diff
} else if max == g_f {
60.0 * (((b_f - r_f) / delta) + 2.0)
} else {
60.0 * (((r_f - g_f) / delta) + 4.0)
};
// OpenCV 的 H 量化H / 2
// 注意OpenCV 底层使用截断还是四舍五入与特定版本有关,
// 标准的 cvtColor 内部实现通常是: h * (180.0 / 360.0) -> h * 0.5
// 这里使用强转(截断)若单测对齐发现差1可改为 (h * 0.5 + 0.5) 或 round()
let h_opencv = (h * 0.5) as u8;
// 3. 计算 S (饱和度)
// OpenCV 公式: S = max == 0 ? 0 : 255 * delta / max
let s_opencv = if max == 0.0 {
0
} else {
((255.0 * delta) / max) as u8
};
// 4. 计算 V (明度)
let v_opencv = max as u8;
(h_opencv, s_opencv, v_opencv)
}