diff --git a/src/lib.rs b/src/lib.rs index cf3276a..d5c68c8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -98,12 +98,13 @@ impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { // 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().color_filter(&ColorPreset::Custom(vec![ - // 错误：下界 (82, 221, 14) 没问题 - // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ - HsvRange::new((82, 221, 14), (240, 203, 82)), - ])).predict(img), + // Runtime::Ocr(s) => s.builder().color_filter(&ColorPreset::Custom(vec![ + // // 错误：下界 (82, 221, 14) 没问题 + // // 但上界的 H 通道写成了 240，超过了 180 的法定上限！ + // HsvRange::new((82, 221, 14), (240, 203, 82)), + // ])).predict(img), Runtime::Det(_) => Err(anyhow::anyhow!("当前模型是检测模型，无法执行 OCR")), } } diff --git a/src/models/ocr.rs b/src/models/ocr.rs index fdc8691..f76279a 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -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 { - 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 { // 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(¤t_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 { diff --git a/src/utils/color_filter.rs b/src/utils/color_filter.rs index e35f389..3237ca2 100644 --- a/src/utils/color_filter.rs +++ b/src/utils/color_filter.rs @@ -1,4 +1,7 @@ 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)] pub struct HsvRange { @@ -6,6 +9,47 @@ pub struct HsvRange { 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 { + // 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::, Vec>::from_raw(width, height, raw_pixels) + .ok_or_else(|| anyhow!("图像缓冲重新组装失败，维度与数据大小不匹配"))?; + + Ok(DynamicImage::ImageRgb8(filtered_buffer)) +} impl HsvRange { pub const fn new(lower: (u8, u8, u8), upper: (u8, u8, u8)) -> Self { Self { lower, upper } @@ -119,6 +163,38 @@ pub trait ColorFilter { fn validate_self(&self) -> Result<(), String> { Ok(()) } + /// 【新扩展的架构方法】将自身安全的合并到已有的普通容器中，并完成去重和排序 + /// 完美的责任分离：Builder 不再需要关心怎么分配内存、怎么排序去重 + fn merge_to_vec(&self, mut existing: Option>) -> Result>, 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 { @@ -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) -} diff --git a/src/utils/cv_ops.rs b/src/utils/cv_ops.rs index 7e22c55..f4c45a7 100644 --- a/src/utils/cv_ops.rs +++ b/src/utils/cv_ops.rs @@ -1,6 +1,6 @@ -use std::cmp::{max, min}; 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) pub fn abs_diff(a: &ArrayView3, b: &ArrayView3) -> Array3 { @@ -13,7 +13,6 @@ pub fn abs_diff(a: &ArrayView3, b: &ArrayView3) -> Array3 { diff } - /// RGB 到灰度转换 pub fn rgb_to_gray(rgb: ArrayView3) -> Array2 { let (h, w, _) = rgb.dim(); @@ -67,9 +66,16 @@ pub fn find_contours_and_max(labelled: &ImageBuffer, Vec>) -> Opt 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, Vec>,max_label: u32) -> (u32, u32, u32, u32) { +pub fn bounding_rect( + labelled: &ImageBuffer, Vec>, + max_label: u32, +) -> (u32, u32, u32, u32) { // 5. 计算最大区域的边界框 (对应 cv2.boundingRect) let mut min_x = labelled.width(); let mut max_x = 0; @@ -85,7 +91,6 @@ pub fn bounding_rect(labelled: &ImageBuffer, Vec>,max_label: u32) } } - let w = max_x - min_x; let h = max_y - min_y; (min_x, min_y, w, h) @@ -104,4 +109,53 @@ pub fn ndarray_to_luma8(array: ArrayView2) -> ImageBuffer, Vec> } } buffer -} \ No newline at end of file +} +// ===================================================================== +// 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) +}