diff --git a/src/lib.rs b/src/lib.rs index 2781c36..cf3276a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,6 +15,7 @@ use models::det::Det; use models::loader::ModelSession; use models::ocr::Ocr; use crate::model_metadata::ModelMetadata; +use crate::utils::color_filter::{ColorPreset, HsvRange}; pub enum ModelSpec { /// 默认 OCR (使用内置路径) @@ -97,7 +98,12 @@ impl DdddOcr { pub fn classification(&self, img: &DynamicImage) -> Result { match &self.runtime { // Runtime::Ocr(s) => s.predict(img).run(), - 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![ + // 错误：下界 (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 f9e10ba..fdc8691 100644 --- a/src/models/ocr.rs +++ b/src/models/ocr.rs @@ -2,6 +2,7 @@ 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::image_io::png_rgba_white_preprocess; use crate::utils::image_processor::{convert_to_grayscale, resize_image}; use anyhow::Context; @@ -126,7 +127,7 @@ pub struct OcrBuilder<'a> { #[allow(dead_code)] probability: bool, /// 颜色过滤：保留的颜色列表 - color_filter_colors: Option>, + color_filter: Result>, String>, /// 颜色过滤：自定义RGB范围 color_filter_custom_ranges: Option>, /// 字符集范围 @@ -141,7 +142,7 @@ impl<'a> OcrBuilder<'a> { // image, png_fix: false, // 默认值 probability: false, - color_filter_colors: None, + color_filter: Ok(None), color_filter_custom_ranges: None, charset_restrict: None, } @@ -150,10 +151,58 @@ impl<'a> OcrBuilder<'a> { self.png_fix = value; self } - pub fn color_filter_colors(mut self, value: Vec) -> Self { - self.color_filter_colors = Some(value); + + // 反复调用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 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 } + pub fn color_filter_custom_ranges(mut self, value: Vec) -> Self { self.color_filter_custom_ranges = Some(value); self diff --git a/src/utils/color_filter.rs b/src/utils/color_filter.rs new file mode 100644 index 0000000..e35f389 --- /dev/null +++ b/src/utils/color_filter.rs @@ -0,0 +1,222 @@ +use std::str::FromStr; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct HsvRange { + pub lower: (u8, u8, u8), // (H, S, V) + pub upper: (u8, u8, u8), // (H, S, V) +} + +impl HsvRange { + pub const fn new(lower: (u8, u8, u8), upper: (u8, u8, u8)) -> Self { + Self { lower, upper } + } +} +impl HsvRange { + /// 验证当前 HSV 范围是否合法 + /// 对应 Python 逻辑：H 在 0-180，S/V 在 0-255，且下界 <= 上界 + pub fn validate(&self) -> Result<(), String> { + // 1. 校验 H 通道边界 (OpenCV 中 H 范围是 0-180) + if self.lower.0 > 180 || self.upper.0 > 180 { + return Err("H通道值必须在 0-180 范围内".to_string()); + } + + // 2. 校验下界不能大于上界 + if self.lower.0 > self.upper.0 || self.lower.1 > self.upper.1 || self.lower.2 > self.upper.2 { + return Err("HSV范围下界不能大于上界".to_string()); + } + + Ok(()) + } +} +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ColorPreset { + Red, + Blue, + Green, + Yellow, + Orange, + Purple, + Cyan, + Black, + White, + Gray, + Custom(Vec), +} + +impl ColorPreset { + /// 纯裸数据定义，没有任何结构体包装，干净利落 + /// 返回值：(范围数量, 范围数组) + /// 完美的零成本抽象：利用常量提升将数据直接打入只读数据段 (.rodata) + pub fn matches(&self) -> &[HsvRange] { + match self { + ColorPreset::Red => &[ + HsvRange { lower: (0, 50, 50), upper: (10, 255, 255) }, + HsvRange { lower: (170, 50, 50), upper: (180, 255, 255) }, + ], + ColorPreset::Blue => &[HsvRange { lower: (100, 50, 50), upper: (130, 255, 255) }], + ColorPreset::Green => &[HsvRange { lower: (40, 50, 50), upper: (80, 255, 255) }], + ColorPreset::Yellow => &[HsvRange { lower: (20, 50, 50), upper: (40, 255, 255) }], + ColorPreset::Orange => &[HsvRange { lower: (10, 50, 50), upper: (20, 255, 255) }], + ColorPreset::Purple => &[HsvRange { lower: (130, 50, 50), upper: (170, 255, 255) }], + ColorPreset::Cyan => &[HsvRange { lower: (80, 50, 50), upper: (100, 255, 255) }], + ColorPreset::Black => &[HsvRange { lower: (0, 0, 0), upper: (180, 255, 50) }], + ColorPreset::White => &[HsvRange { lower: (0, 0, 200), upper: (180, 30, 255) }], + ColorPreset::Gray => &[HsvRange { lower: (0, 0, 50), upper: (180, 30, 200) }], + ColorPreset::Custom(ranges) => ranges, + } + } + /// 校验逻辑：在这里实现完美的“责任分离” + pub fn validate(&self) -> Result<(), String> { + match self { + // 1. 快捷变体：完全绕过，根本不校验，0 运行时开销放行！ + ColorPreset::Custom(ranges) => { + // 2. 只有 Custom 变体需要接受严格的参数政审 + for r in ranges { + r.validate()?; + } + Ok(()) + } + _ => Ok(()), + } + } +} + +impl FromStr for ColorPreset { + type Err = String; + fn from_str(s: &str) -> Result { + match s.to_lowercase().as_str() { + "red" => Ok(ColorPreset::Red), + "blue" => Ok(ColorPreset::Blue), + "green" => Ok(ColorPreset::Green), + "yellow" => Ok(ColorPreset::Yellow), + "orange" => Ok(ColorPreset::Orange), + "purple" => Ok(ColorPreset::Purple), + "cyan" => Ok(ColorPreset::Cyan), + "black" => Ok(ColorPreset::Black), + "white" => Ok(ColorPreset::White), + "gray" => Ok(ColorPreset::Gray), + _ => Err(format!("不支持的颜色预设: {}", s)), + } + } +} + +// ===================================================================== +// 3. 颜色约束特征（Trait）与组合子设计模式 +// ===================================================================== + +pub struct PixelCtx { + pub hsv: (u8, u8, u8), +} + +/// 统一的颜色约束接口 +pub trait ColorFilter { + /// 将自身的有效约束平铺追加到统一的目标容器中 + fn append_ranges(&self, target: &mut Vec); + /// 预估范围数量，借助原生内置的 len() 实现 O(1) 完美控容 + fn estimated_count(&self) -> usize; + /// 将自身的有效约束平铺追加到统一目标容器中 + /// 验证当前过滤器是否合法，默认直接放行（Ok(())） + fn validate_self(&self) -> Result<(), String> { + Ok(()) + } +} + +impl ColorFilter for ColorPreset { + fn append_ranges(&self, target: &mut Vec) { + // 直接利用我们第一步写好的 matches() 拿到切片，整块高速拷贝倒入目标容器 + target.extend_from_slice(self.matches()); + } + fn estimated_count(&self) -> usize { + // 直接获取切片长度 + self.matches().len() + } + fn validate_self(&self) -> Result<(), String> { + // 直接调用我们在第一步中为 ColorPreset 实现的精细化分流校验 + // 快捷变体在这里会直接返回 Ok(()), 只有 Custom 才会去真正校验 + self.validate() + } +} + + +/// 多路颜色“或”逻辑组合子（并集网络） +pub struct MultiOrColorRestrict<'a> { + pub filters: Vec<&'a dyn ColorFilter>, +} + +impl<'a> ColorFilter for MultiOrColorRestrict<'a> { + fn append_ranges(&self, target: &mut Vec) { + // 管道递延：依次指挥内部每一个子过滤器把数据倒进目标容器 + for f in &self.filters { + f.append_ranges(target); + } + } + fn estimated_count(&self) -> usize { + // 数量累加：$O(1)$ 地把所有子过滤器的预估容量加起来 + self.filters.iter().map(|f| f.estimated_count()).sum() + } + + fn validate_self(&self) -> Result<(), String> { + // 递归政审：只要其中一个子过滤器校验失败（比如某个 Custom 变体非法），立刻熔断 + for f in &self.filters { + f.validate_self()?; + } + Ok(()) + } +} + +// ===================================================================== +// 4. 声明式宏：一语定乾坤 +// ===================================================================== + +#[macro_export] +macro_rules! color_any_of { + ($only:expr) => { + &$only as &dyn $crate::ColorFilter + }; + ($($filter:expr),+ $(,)?) => { + &$crate::MultiOrColorRestrict { + filters: vec![ $( &$filter as &dyn $crate::ColorFilter ),+ ] + } + }; +} + +// ===================================================================== +// 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/mod.rs b/src/utils/mod.rs index 0d04ccf..7e0f806 100644 --- a/src/utils/mod.rs +++ b/src/utils/mod.rs @@ -1,3 +1,4 @@ pub mod image_io; pub mod image_processor; -pub mod cv_ops; \ No newline at end of file +pub mod cv_ops; +pub mod color_filter; \ No newline at end of file diff --git a/tests/ocr_test.rs b/tests/ocr_test.rs index d6db696..928eee6 100644 --- a/tests/ocr_test.rs +++ b/tests/ocr_test.rs @@ -66,7 +66,7 @@ fn test_full_classification() { let ocr = DdddOcrBuilder::new().build().expect("模型加载失败"); // 2. 加载测试图片 - let img = image::open("samples/code3.png").expect("测试图片不存在"); + let img = image::open("samples/code2.png").expect("测试图片不存在"); // 3. 执行识别 let result = ocr.classification(&img).expect("识别过程出错");