From 48c2cbedb0b4065297828962e85ca0b801bac468 Mon Sep 17 00:00:00 2001 From: CNWei Date: Mon, 29 Jun 2026 17:19:01 +0800 Subject: [PATCH] =?UTF-8?q?refactor(ocr):=20=E4=BC=98=E5=8C=96=E5=B9=B6?= =?UTF-8?q?=E7=B2=BE=E7=AE=80=20color=5Ffilter=20=E6=9E=B6=E6=9E=84?= =?UTF-8?q?=E8=AE=BE=E8=AE=A1?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 重构 `OcrBuilder`，将图像矩阵过滤与像素比对等执行层逻辑彻底剥离解耦。 - 优化 `OcrBuilder` 的 `color_filter` 链式调用，将其改造为无心智负担的单次覆盖（Overwrite）逻辑。 - 扩展 `ColorFilter` 特征，新增 `collect_to_vec` 方法，实现底层规则的高内聚收集、精准内存开辟与原地去重排序。 --- src/charset.rs | 43 +++++++++++++++++ src/models/ocr.rs | 97 ++++++--------------------------------- src/utils/color_filter.rs | 17 ++----- 3 files changed, 60 insertions(+), 97 deletions(-) diff --git a/src/charset.rs b/src/charset.rs index 6a88a71..c9de064 100644 --- a/src/charset.rs +++ b/src/charset.rs @@ -533,6 +533,49 @@ pub trait TokenFilter { fn estimated_capacity(&self) -> usize { 128 } + /// 【新引入的架构级核心方法】 + /// 统一接管全量字符集的密集遍历、CTC Blank放行、去重、排序及空交集退化兜底 + fn apply_to_charset(&self, tokens: &[Cow]) -> Option> { + let mut has_any_match = false; + let estimated_capacity = self.estimated_capacity(); + + // 1. 精准开辟内存，完美利用容量提示，避免动态乱涨 + let mut temp_indices = Vec::with_capacity(estimated_capacity.max(16)); + + // 2. 高性能原地单次流式迭代 + for (idx, token) in tokens.iter().enumerate() { + let token_str = token.as_ref(); + + // 规则 A: CTC Blank 空字符串或 0 号索引无条件放行 + if token_str.is_empty() || idx == 0 { + temp_indices.push(idx); + continue; // 关键：直接跳过，防止后续 matches 匹配成功导致重复 push 产生 Bug + } + + // 规则 B: 组装无拷贝上下文 + let ctx = ValidationCtx { + text: token_str, + token_id: idx, + }; + + // 规则 C: 路由到各自具体实现的特异性匹配中（如 Digit 判定、TopN 判定、组合子判定等） + if self.matches(&ctx) { + temp_indices.push(idx); + has_any_match = true; + } + } + + // 3. 终极防御：如果整个模型字符集除了 Blank，一个都没对上，直接退化为 None（全量识别） + if !has_any_match { + println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); + None + } else { + // 4. 排序并去重，为 Ocr 引擎后续进行极其高频的『二分查找』筑起绝对安全的底层保障 + temp_indices.sort_unstable(); + temp_indices.dedup(); + Some(temp_indices) + } + } } #[derive(Debug, Clone, PartialEq, Eq)] diff --git a/src/models/ocr.rs b/src/models/ocr.rs index f76279a..b2ac19d 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -1,14 +1,14 @@ -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::{filter_image, ColorFilter, HsvRange}; +use crate::utils::color_filter::{ColorFilter, HsvRange, filter_image}; use crate::utils::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use anyhow::Context; -use anyhow::{anyhow, Result}; +use anyhow::{Result, anyhow}; use image::{DynamicImage, ImageBuffer, Rgb}; +use std::borrow::Cow; use std::collections::HashSet; use tract_onnx::prelude::tract_ndarray::s; use tract_onnx::prelude::{ @@ -158,48 +158,13 @@ impl<'a> OcrBuilder<'a> { // 反复调用color_filter怎么处理？ pub fn color_filter(mut self, filter: &dyn ColorFilter) -> Self { // self.color_filter = Some(value); - // 利用组合子预估能力，获取精准分配槽位数 - let total_capacity = filter.estimated_count(); - if total_capacity == 0 { - return self; + + // 一句话把活全包了！错误信息无缝传递，完美熔断 + match filter.collect_to_vec() { + Ok(new_ranges) => self.color_filter = Ok(new_ranges), + Err(err_msg) => self.color_filter = Err(err_msg), // 校验失败，Builder 正式中毒 } - match self.color_filter { - Ok(mut ranges) => { - // 2. 触发特征多态自检：快捷预设秒过，Custom 变体严格政审 - if let Err(err_msg) = filter.validate_self() { - // 校验失败，Builder 正式中毒，熔断器闭合 - self.color_filter = Err(err_msg); - return self; - } - - match &mut ranges { - None => { - // 情况 A：这是用户第一次配置，直接一击必中分配精准内存 - let mut v = Vec::with_capacity(total_capacity); - - // 尝试倒入，如果出错，Builder 正式中毒 - filter.append_ranges(&mut v); - ranges = Some(v); - } - Some(existing_vec) => { - // 情况 B：用户之前配置过，现在追加配置。 - // 绝不让它动态乱涨！我们利用 reserve 方法，【仅此一次】精准追加所需的扩容空间！ - existing_vec.reserve(total_capacity); - - // 尝试追加倒入 - filter.append_ranges(existing_vec); - } - } - if let Some(v) = &mut ranges { - v.sort_unstable(); - v.dedup(); - } - // 5. 更新状态 - self.color_filter = Ok(ranges); - } - Err(_) => return self, - }; self } @@ -211,42 +176,7 @@ impl<'a> OcrBuilder<'a> { let charset = &self.ocr.model_metadata.charset; let tokens = &charset.tokens; // let mut temp_indices = Vec::new(); - let mut has_any_match = false; - - let estimated_capacity = restrict.estimated_capacity(); - // 精准开辟内存，完美避开 8210 个槽位的巨大空置浪费 - let mut temp_indices = Vec::with_capacity(estimated_capacity); - - for (idx, token) in tokens.iter().enumerate() { - let token_str = token.as_ref(); - // CTC Blank 空字符串无条件放行，其余交给超高性能的 matches - if token_str.is_empty() || idx == 0 { - temp_indices.push(idx); - }; - // 组装无拷贝上下文 - let ctx = ValidationCtx { - text: token_str, - token_id: idx, - }; - - if restrict.matches(&ctx) { - temp_indices.push(idx); - has_any_match = true; - } - } - - // self.charset_restrict = Some(restrict); - // 终极防御：如果除了 Blank 外什么都没匹配上，退化恢复为 None（全量识别） - if !has_any_match { - println!("警告：当前限制策略与模型字符集完全没有交集！已自动恢复全量识别。"); - self.charset_restrict = None; - } else { - // 这一步非常重要：二分查找（binary_search）强依赖数组【有序】。 - // TopN 天然有序，但如果是用户自定义的 CustomList 或者复杂的 Or 组合， - // 遍历出来的索引天然有序，但为了绝对的安全，我们在这里顺手排个序 - temp_indices.sort_unstable(); - self.charset_restrict = Some(temp_indices); - } + self.charset_restrict = restrict.apply_to_charset(tokens); self } } @@ -260,7 +190,10 @@ impl<'a> OcrBuilder<'a> { // ===================================================================== let img_cow = match &self.color_filter { Err(err_msg) => { - return Err(anyhow::anyhow!("颜色过滤器初始化失败，全链路短路: {}", err_msg)); + return Err(anyhow::anyhow!( + "颜色过滤器初始化失败，全链路短路: {}", + err_msg + )); } Ok(None) => { // 核心优化点：直接借用原图，不发生任何克隆 @@ -313,10 +246,6 @@ 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 3237ca2..24bb808 100644 --- a/src/utils/color_filter.rs +++ b/src/utils/color_filter.rs @@ -165,26 +165,17 @@ pub trait ColorFilter { } /// 【新扩展的架构方法】将自身安全的合并到已有的普通容器中，并完成去重和排序 /// 完美的责任分离：Builder 不再需要关心怎么分配内存、怎么排序去重 - fn merge_to_vec(&self, mut existing: Option>) -> Result>, String> { + fn collect_to_vec(&self) -> Result>, String> { // 1. 触发自检 self.validate_self()?; let total_capacity = self.estimated_count(); if total_capacity == 0 { - return Ok(existing); + return Ok(None); } - 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. 永远一击必中分配精准内存，不需要再考虑追加和扩容！ + let mut v = Vec::with_capacity(total_capacity.max(16)); // 2. 倒入数据 self.append_ranges(&mut v);